from transformers import AutoModelForSeq2SeqLM
import torch
# Initialize the tokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Sunbird/sunbird-mul-en-mbart-merged")

# Initialize the model
model = AutoModelForSeq2SeqLM.from_pretrained("Sunbird/sunbird-mul-en-mbart-merged")
model.load_state_dict(torch.load("Total Combined Data V2 Aug 16 2023/Models/mul_en_base.bin"))
model = model.to("cuda")  # if you're using a GPU

sentence = "Kya busiru okuluubirira eby'obugagga, Eddwaliro terisobola kusuza balwadde bangi nnyo. Tekisoboka"
# Tokenize the sentence
tokenized_sentence = tokenizer(sentence, return_tensors="pt")
tokenized_sentence = tokenized_sentence.to("cuda")  # if you're using a GPU
# Generate the translation
translated_tokens = model.generate(**tokenized_sentence)
# Decode the translation
translation = tokenizer.decode(translated_tokens[0])
print("*"*6)
print("\n")
print(translation)
print("\n")
print("*"*6)