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],
        ["mi li toki e toki pona", "toki pona", "toki pona", 3]
    ],
    cache_examples=False,
    title="English / toki pona Translation",
    description="""
    # A simple English / toki pona Neural Machine Translation App!
    
    ### toki a! 💬

    This is a simple english to toki pona / toki pona to english neural machine translation app.

    Input your text to translate, a source language and target language, and desired number of return sequences! 

    ### Grammaticality / Regularization
    English -> English and/or toki pona -> toki pona will result in some form of regularization. 
    
    This can approximate grammaticality, but it isn't always the best.
    
    For example, "mi li toki e toki pona" [src: toki pona, tgt: toki pona] will result in ['mi toki e toki pona.', 'mi toki pona.', 'mi toki e toki pona']
    (Thus, the ungrammatical "li" is dropped)

    ### Model and Data
    This app utilizes a fine-tuned version of Facebook/Meta AI's M2M100 418M param model. 
    
    By leveraging the pretrained weights of the massively multilingual M2M100 model, 
    we can jumpstart our transfer learning to accomplish machine translation for toki pona!
    
    The model was fine-tuned on the English/toki pona bitexts found at https://tatoeba.org/
    
    ### This app is a work in progress and obviously not all translations will be perfect. 
    In addition to parameter quantity and the hyper-parameters used while training, 
    the *quality of data* found on Tatoeba directly influences the perfomance of projects like this! 

    If you wish to contribute, please simply add high quality and diverse translations to Tatoeba!

    """
)

app.launch()