import os
import torch
from convert_nemo_ul2_checkpoint import convert_nemo_to_hf
from transformers import T5ForConditionalGeneration, AutoTokenizer

#### Step 1: Convert the original HF model which was converted to NEMO back to HF weights
nemo_weights = torch.load("ul2-base-nl36-finnish/nemo_state_dict.pt")
hf_weights = convert_nemo_to_hf(nemo_weights)

#### Step 2: Load original HF model and save its config/tokenizer in local folder
hf_model = T5ForConditionalGeneration.from_pretrained("Finnish-NLP/ul2-base-nl36-finnish")
tokenizer = AutoTokenizer.from_pretrained("Finnish-NLP/ul2-base-nl36-finnish")

# Save tokenizer in ul2-base-nl36-finnish
tokenizer.save_pretrained("ul2-base-nl36-finnish/hf_t5_ul2")

# Save config in ul2-base-nl36-finnish
hf_model.config.save_pretrained("ul2-base-nl36-finnish/hf_t5_ul2")

#### Step 3: Save our converted weights to the local folder
# Save converted model weights in ul2-base-nl36-finnish
torch.save(hf_weights, os.path.join("ul2-base-nl36-finnish/hf_t5_ul2", "pytorch_model.bin"))


#### Step4: Load the converted model from local folder and check whether weights are the same
converted_model = T5ForConditionalGeneration.from_pretrained("ul2-base-nl36-finnish/hf_t5_ul2")

equal = []
for key in hf_model.state_dict().keys():
    print(key)
    print(torch.allclose(hf_model.state_dict()[key], converted_model.state_dict()[key]))
    equal.append(torch.allclose(hf_model.state_dict()[key], converted_model.state_dict()[key]))


print(f"All weights are equal: {all(equal)}")