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