import gradio as gr

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

modelcard = "amurienne/gallek-m2m100"

model = AutoModelForSeq2SeqLM.from_pretrained(modelcard)
tokenizer = AutoTokenizer.from_pretrained(modelcard)

def translate(text):
    """
    Translate the text from source lang fr to target lang br
    """
    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, src_lang='fr', tgt_lang='br', max_length=400, device="cpu")
    result = translation_pipeline("traduis de français en breton: " + text)
    return result[0]['translation_text']

demo = gr.Interface(
    fn=translate,
    inputs=[
        gr.components.Textbox(label="French"),
    ],
    outputs=[
        gr.components.Textbox(label="Breton")
    ],
    cache_examples=False,
    title="Gallek French -> Breton Translation Demo",
)

demo.launch()