import gradio as gr
from transformers import MarianMTModel, MarianTokenizer


def gradio_input(txt):
    translated = translate(txt)
    return translated[0]


def translate(text):
    hub_repo_name = 'dataequity/dataequity-kde4-en-es-qlora'
    tokenizer = MarianTokenizer.from_pretrained(hub_repo_name)
    finetuned_model = MarianMTModel.from_pretrained(hub_repo_name)

    try:
        translated = finetuned_model.generate(**tokenizer([text], return_tensors="pt", padding=True))
        return [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
    except:
        return ['Invalid input']


iface = gr.Interface(fn=gradio_input,
                     inputs=gr.Textbox(placeholder="English text"),
                     outputs=gr.Textbox(placeholder="Translation"),
                     title="Translate English to Spanish")
iface.launch()