import gradio as gr

def greet(name):
    return "Hello " + name + "!!"


def translate(input):
    from transformers import T5ForConditionalGeneration, T5Tokenizer

    model_name = 'jbochi/madlad400-3b-mt'
    model = T5ForConditionalGeneration.from_pretrained(model_name, device_map="auto")
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    
    text = f"<2en>{input}"
    input_ids = tokenizer(text, return_tensors="pt").input_ids.to(model.device)
    outputs = model.generate(input_ids=input_ids)
    
    output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Eu adoro pizza!

    return output

demo = gr.Interface(fn=translate, inputs="text", outputs="text", title = 'Latin Translator', examples=[['Mi amiga Morgan es bonita'],['Hola, buenos noches!']])
demo.launch(share=True)