File size: 2,427 Bytes
0fd282e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration, T5Tokenizer
from nemo.collections.nlp.models.language_modeling.megatron_t5_model import MegatronT5Model
from nemo.collections.nlp.data.language_modeling.megatron.ul2_dataset import UL2Dataset
from pytorch_lightning.trainer.trainer import Trainer
def load_nemo_megatron_model(checkpoint_path, devices=1, num_nodes=1, accelerator="gpu"):
trainer = Trainer(devices=devices, num_nodes=num_nodes, accelerator=accelerator)
model = MegatronT5Model.load_from_checkpoint(checkpoint_path, trainer=trainer)
return model
#### Huggingface ####
tokenizer = AutoTokenizer.from_pretrained("ul2-base-nl36-swedish")
model = T5ForConditionalGeneration.from_pretrained("ul2-base-nl36-swedish")
# "Hunden bet mannen i" means "The dog bit the man in".
input_ids = tokenizer(
"<extra_id_r> Hunden bet mannen i <extra_id_0>", return_tensors="pt", return_token_type_ids=False
)
# Predict with HF
with torch.no_grad():
outputs_hf = model(
input_ids=input_ids.input_ids,
attention_mask=input_ids.attention_mask,
decoder_input_ids=input_ids.input_ids,
decoder_attention_mask=input_ids.attention_mask,
)
# Argmax to get the most probable token id
output_tokens_hf = outputs_hf[0].argmax(dim=-1)
#### Nemo ####
model_nemo = load_nemo_megatron_model("nemo_checkpoints/megatron_ul2--val_loss=2.54-step=7000-consumed_samples=14557920.0.ckpt")
model_nemo.eval()
tokenizer_nemo = model_nemo.tokenizer.tokenizer
input_ids_nemo = tokenizer_nemo("<extra_id_r> Hunden bet mannen i <extra_id_0>", return_tensors="pt").to("cuda")
# Predict with Nemo
with torch.no_grad():
outputs_nemo = model_nemo(
encoder_input_ids=input_ids_nemo.input_ids,
decoder_input_ids=input_ids_nemo.input_ids,
encoder_attn_mask=input_ids_nemo.attention_mask,
decoder_attn_mask=input_ids_nemo.attention_mask,
)
# Argmax to get the most probable token
output_tokens = outputs_nemo.argmax(dim=-1)
#### Compare both outputs ####
print(f"Nemo logits: {outputs_nemo[0]}")
print(f"Huggingface logits: {outputs_hf[0]}")
print(f"Are logits equal: {torch.allclose(outputs_nemo[0], outputs_hf[0].to('cuda'))}")
# Decode tokens
print(f"Huggingface output: {tokenizer.batch_decode(output_tokens_hf)}")
print(f"Nemo output: {tokenizer_nemo.batch_decode(output_tokens)}") # Reasonable output for undertrained model |