from transformers import AutoTokenizer, MT5ForConditionalGeneration
import gradio as gr

model_name = "alfsnd/mt5-base-spanish-yoremnokki"

tokenizer = AutoTokenizer.from_pretrained(model_name)

model = MT5ForConditionalGeneration.from_pretrained(model_name)

def translate(input_text, src_lang, to_lang):
    prompt = f"translate {src_lang} to {to_lang}: {input_text}"
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    outputs = model.generate(input_ids, max_new_tokens=1000)
    model_translation = tokenizer.decode(outputs[0])
    final_translation = model_translation[5:-4]
    return final_translation

languages = [
    'Spanish', 'Yoremnokki'
]

desc = "<p>Traductor de Español a Yoremnokki utilizando un modelo de inteligecia artifical (alfsnd/mt5-base-spanish-yoremnokki)</p>"

translator = gr.Interface(fn=translate,
                          inputs=[gr.Textbox(label="Traducir", placeholder="Introduzca su texto"),
                                  gr.Dropdown(label="De",
                                              choices=languages,
                                              value="Spanish",),
                                  gr.Dropdown(label="A",
                                              choices=languages,
                                              value="Yoremnokki")],
                          outputs=gr.Textbox(label="Traducción"),
                          title="Traductor de Español a Yoremnokki",
                          description=desc,
                          submit_btn="Traducir",
                          clear_btn="Limpiar",
                          )

translator.launch()