import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import torch

# this model was loaded from https://hf.co/models
model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona")
tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
device = "cuda:0" if torch.cuda.is_available() else "cpu"
LANG_CODES = {
    "English":"en",
    "toki pona":"tl"
}

def translate(text, src_lang, tgt_lang, candidates:int):
    """
    Translate the text from source lang to target lang
    """

    src = LANG_CODES.get(src_lang)
    tgt = LANG_CODES.get(tgt_lang)

    tokenizer.src_lang = src
    tokenizer.tgt_lang = tgt

    ins = tokenizer(text, return_tensors='pt').to(device)

    gen_args = {
            'return_dict_in_generate': True,
            'output_scores': True,
            'output_hidden_states': True,
            'length_penalty': 0.0,  # don't encourage longer or shorter output,
            'num_return_sequences': candidates,
            'num_beams':candidates,
            'forced_bos_token_id': tokenizer.lang_code_to_id[tgt]
        }
    

    outs = model.generate(**{**ins, **gen_args})
    output = tokenizer.batch_decode(outs.sequences, skip_special_tokens=True)

    return output

app = gr.Interface(
    fn=translate,
    inputs=[
        gr.components.Textbox(label="Text"),
        gr.components.Dropdown(label="Source Language", choices=list(LANG_CODES.keys())),
        gr.components.Dropdown(label="Target Language", choices=list(LANG_CODES.keys())),
        gr.Slider(label="Number of return sequences", value=3, minimum=1, maximum=12, step=1)
    ],
    outputs=["text"],
    examples=[
        ["Welcome to my translation app.", "English", "toki pona", 3],
        ["Its not always perfect, but its pretty okay!", "English", "toki pona", 3],
        ["ilo pi ante toki ni li pona a!", "toki pona", "English", 3],
        ["kijetesantakalu li pona", "toki pona", "English", 3],
    ],
    cache_examples=False,
    title="A simple English / toki pona Neural Translation App",
    description="A simple English / toki pona Neural Translation App"
)

app.launch()