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

# languages and model
LANGS = ["eng_Latn", "fra_Latn", "spa_Latn"]
model_name = "facebook/nllb-200-distilled-600M"

# model and tokinizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
device = 0 if torch.cuda.is_available() else -1

# translate function
def translate(input, src_lang, tgt_lang):
  translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer, 
                                  src_lang=src_lang, tgt_lang=tgt_lang,
                                  max_length=400, device=device)
  res = translation_pipeline(input)
  return res[0]['translation_text']
##
##
with gr.Blocks() as demo:
    with gr.Row():
      gr.Markdown("""
                  > 👽 0xZee
                  # 📝 NLP : Translation Demo
                  > Demo #1 : **facebook/nllb-200-distilled-600M** Model (finetuned) 
                  *Translation : English, French, Spanish*
                  """)
#
    with gr.Row():
          input = gr.Textbox(label="Text to translate", placeholder="Text here")
    with gr.Row(): 
          src = gr.Dropdown(label="Source Language", choices=LANGS)
          tgt = gr.Dropdown(label="Target Language", choices=LANGS)
    with gr.Row():
          gr.Examples(examples=[["Building a better world, with sustainable energies", "eng_Latn", "spa_Latn"], ["Messi es el mejor del mundo, y de lejos", "spa_Latn", "fra_Latn"]], inputs=[input, src, tgt])
    with gr.Row():
          btn = gr.Button(value="Go Translate")  
    with gr.Row():      
          text_out = gr.Textbox(label="Translation")
          btn.click(fn=translate, inputs=[input, src, tgt], outputs=text_out)



if __name__ == "__main__":
  demo.launch(debug=True)