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

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",
    "Zelsik":"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 '\n'.join(output)

with gr.Blocks() as app:
    markdown="""
    # An English / Zelsik Neural Machine Translation App!
    
    This is an English to Zelsik / Zelsik to English neural machine translation app.

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

    ### Grammar Regularization
    An interesting quirk of training a many-to-many translation model is that pseudo-grammar correction 
    can be achieved by translating *from* **language A** *to* **language A**
    
    Remember, this can ***approximate*** grammaticality, but it isn't always the best.
    

    ### Model and Data
    This app utilizes a fine-tuned version of Facebook/Meta AI's M2M100 418M param model and the original app was made by Jayyydyyy for Toki Pona.
    
    By leveraging the pre-trained weights of the massively multilingual M2M100 model, 
    we can jumpstart our transfer learning to accomplish machine translation for Zelsik!
    
    The model was fine-tuned on the English/Zelsik bitexts found at [https://tatoeba.org/](https://tatoeba.org/)
    
    ### This app is a work in progress; 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 performance of projects like this! 

    im sorry jayyydyyy, im too lazy and dumb to change any of the descriptions
    """

    with gr.Row():
        gr.Markdown(markdown)
        with gr.Column():
            input_text = gr.components.Textbox(label="Input Text", value="Raccoons are fascinating creatures, but I prefer opossums.")
            source_lang = gr.components.Dropdown(label="Source Language", value="English", choices=list(LANG_CODES.keys()))
            target_lang = gr.components.Dropdown(label="Target Language", value="Zelsik", choices=list(LANG_CODES.keys()))
            return_seqs = gr.Slider(label="Number of return sequences", value=3, minimum=1, maximum=12, step=1)
            
            inputs=[input_text, source_lang, target_lang, return_seqs]
            outputs = gr.Textbox()

            translate_btn = gr.Button("Translate! | o ante toki!")
            translate_btn.click(translate, inputs=inputs, outputs=outputs)

            gr.Examples(
                [
                    ["Hello! How are you?", "English", "Zelsik", 3],
                    ["toki a! ilo pi ante toki ni li pona!", "Zelsik", "English",  3],
                    ["mi li toki e Zelsik", "Zelsik", "Zelsik", 3],
                ],
                inputs=inputs
            )

app.launch()