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

model = AutoModelForSeq2SeqLM.from_pretrained("Mamadou2727/Feriji_model")
tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
device = "cuda:0" if torch.cuda.is_available() else "cpu"
LANG_CODES = {
    "French": "fr",
    "Zarma": "yo"
}

def translate(text, candidates: int):
    """
    Translate the text from French to Zarma
    """

    src = LANG_CODES["French"]
    tgt = LANG_CODES["Zarma"]
    
    # If a file is uploaded, use its content for translation
    if file is not None:
        text = file.read().decode("utf-8")
        
    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 = r"""
        # Feriji-fr-to-dje v.1.1, Proudly made by Elysabhete, Habibatou & Mamadou K.
        
        <img src="https://cdn-uploads.huggingface.co/production/uploads/63cc1d4bf488db9bb3c6449e/AtOKLAaL5kt0VhRsxE0vf.png" width="500" height="300">
        
        Feriji-fr-to-dje is a beta version of the French to Zarma translator.
            
        ## Intended Uses & Limitations        
        
        This model is intended for academic research and practical applications in machine translation. It can be used to translate French text to Zarma and vice versa. Users should note that the model's performance may vary based on the complexity and context of the input text.
            
        ## Authors:
        The project, **Feriji dataset and Feriji-fr-to-dje**, was curated by **Elysabhete Ibrahim Amadou** and **Mamadou K. KEITA**, with the aim to enhance linguistic studies and translation capabilities between French and Zarma.
            
        ## Citations
            
        If you use this dataset or model in your research, please cite it as follows:
            
        @dataset{Feriji,
          author       = {Habibatou Abdoulaye Alfari, Elysabhete Ibrahim Amadou and Mamadou K. KEITA},
          title        = {Feriji, a French-Zarma Parallel Corpus},
          year         = 2023,
          publisher    = {GitHub},
          journal      = {GitHub repository},
          howpublished = {\url{https://github.com/27-GROUP/Feriji}}
        }
    """

    with gr.Row():
        gr.Markdown(markdown)
        with gr.Column():
            input_text = gr.components.Textbox(lines=7, label="Input Text", value="")
            upload_file = gr.File(label="Upload File")
            return_seqs = gr.Slider(label="Number of return sequences", value=1, minimum=1, maximum=12, step=1)
            outputs = gr.Textbox(lines=7, label="Output Text")

            translate_btn = gr.Button("Traduis!")
            translate_btn.click(translate, inputs=[input_text, upload_file, return_seqs], outputs=outputs)
            
app.launch(share=True)