from multilingual_translation import text_to_text_generation
from utils import data_scraping
import gradio as gr

model_list = data_scraping()

def multilingual_translate(prompt: str, model_id: str):
    target_lang = "en" # English language code
    return text_to_text_generation(
        prompt=prompt, 
        model_id=model_id, 
        device='cpu',
        target_lang=target_lang
    )

inputs = [
    gr.Textbox(lines=4, value="Hello world!", label="Input Text"),
    gr.Dropdown(model_list, value="facebook/m2m100_418M", label="Model"),
]

output = gr.outputs.Textbox(label="Output Text")

examples = [
    ["Hello world!", "facebook/m2m100_418M"],
    ["Omar ve Merve çok iyi arkadaşlar.", "facebook/m2m100_418M"],
    ["Hugging Face is a great company.", "facebook/m2m100_418M"]
]

title = "Beyond English-Centric Multilingual Machine Translation"
description = "M2M100 is a multilingual encoder-decoder (seq-to-seq) model trained for Many-to-Many multilingual translation. It was introduced in this [paper](https://arxiv.org/abs/2010.11125) and first released in [this](https://github.com/pytorch/fairseq/tree/master/examples/m2m_100) repository."

app = gr.Interface(
    fn=multilingual_translate,
    inputs=inputs,
    outputs=output,
    examples=examples,
    title=title,
    description=description,
    cache_examples=True
)
app.launch(debug=True, enable_queue=True)