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

REPO_ID_NLLB = "facebook/nllb-200-distilled-600M"
REPO_ID_MARIANNMT_en = "mbarnig/MarianNMT-tatoeba-en-lb"
REPO_ID_MARIANNMT_lb = "mbarnig/MarianNMT-tatoeba-lb-en"
REPO_ID_T5MT5 = "mbarnig/T5-mt5-tatoeba-en-lb"

my_title = "🇬🇧 Mir iwwersetzen vun an op Lëtzebuergesch ! 🇫🇷"
my_description = "English-Luxembourgish machine translation (MT) demo based on 3 open-source transformer models: Facebook-NLLB, Microsoft-MarianNMT & Google-T5/mt5."
my_article = "<h3>User guide</h3><p>1. Press the submit button to translate an english text with the default values. 2. Compare the result with the luxembourgish example. 3. Select a model and a translation direction and enter your own text. Have fun !</p><p>Go to <a href='https://www.web3.lu/'>Internet with a Brain</a> to read my french publication <a href='https://www.web3.lu/'>Das Küsschen und die Sonne stritten sich ...</a> about the history of machine translation in Luxembourg from 1975 until today.</p>"
default_input = "The North Wind and the Sun were disputing which was the stronger, when a traveler came along wrapped in a warm cloak."

TRANSLATION_MODELS = [
    "NLLB",
    "MarianNMT",
    "T5/mt5"
]

TRANSLATION_DIRECTION = [
    "en -> lb",
    "lb -> en"
]

EXAMPLE = "..."

my_inputs = [
    gr.Textbox(lines=5, label="Input", value=default_input),
    gr.Radio(label="Translation Model", choices = TRANSLATION_MODELS, value = "NLLB"),
    gr.Radio(label="Translation Direction", choices = TRANSLATION_DIRECTION, value = "en -> lb")
]

my_output = gr.Textbox(lines=5, label="Translation")

def customization(myModel, direc):
    if myModel == "NLLB":
        translator = pipeline("translation", model=REPO_ID_NLLB)   
    elif myModel == "MarianNMT":
        if direc == "en -> lb":
            translator = pipeline("translation", model=REPO_ID_MARIANNMT_en)        
        elif direc == "lb -> en":
            translator = pipeline("translation", model=REPO_ID_MARIANNMT_lb)        
        else:
            print("Please select a Translation Direction !")    
    elif myModel == "T5/mt5":
       translator = pipeline("translation", model=REPO_ID_T5MT5) 
       # tokenizer = AutoTokenizer.from_pretrained("mbarnig/T5-mt5-taboeta-en-lb")
       # model = AutoModelForSeq2SeqLM.from_pretrained("mbarnig/T5-mt5-taboeta-en-lb")       
    else:
       print("Please select a Translation Model !")
    return myModel

def iwwersetz(source_text, model, direc):
    translator = customization(model, direc)
    if model == "NLLB":
        if direc == "en -> lb":
            # translation = translator("en", "lb", source_text)
            translation = source_text
        else:
            # translation = translator("lb", "en", source_text) 
            translation = source_text       
    else:
        # translation = translator(source_text)
        translation = source_text
    return translation
    
demo=gr.Interface(
   fn=iwwersetz,
   inputs=my_inputs,
   outputs=my_output,
   title=my_title, 
   description=my_description, 
   article=my_article,
   allow_flagging=False)
demo.launch()