from transformers import MT5ForConditionalGeneration, AutoTokenizer
from peft import PeftModel

model = MT5ForConditionalGeneration.from_pretrained('google/mt5-large')
tokennizer = AutoTokenizer.from_pretrained('google/mt5-large')

model = PeftModel.from_pretrained(model, 'WilliamSotoM/PTHQL_level0_Indo_European')
model = model.merge_and_unload()

model = PeftModel.from_pretrained(model, 'WilliamSotoM/PTHQL_level1_Germanic')
model = model.merge_and_unload()

model = PeftModel.from_pretrained(model, 'WilliamSotoM/PTHQL_level2_High_German')
model = model.merge_and_unload()

model = PeftModel.from_pretrained(model, 'WilliamSotoM/PTHQL_language_German')
model = model.merge_and_unload()

graph = '''
(w / want-01
   :ARG0 (b / boy)
   :ARG1 (b2 / believe-01
             :ARG0 (g / girl)
             :ARG1 b))
'''
tokenized_input = tokenizer(graph, return_tensors='pt')

with torch.inference_mode():
    prediction = model.generate(**tokenized_input)
    generated_text = tokenizer.batch_decode(prediction, skip_special_tokens=True)[0]

print(f'Generated text:', generated_text)