'''
from easynmt import EasyNMT
import os
trans_model = EasyNMT("m2m_100_1.2B")
'''

from hf_hub_ctranslate2 import MultiLingualTranslatorCT2fromHfHub
from transformers import AutoTokenizer
import gradio as gr

tokenizer = AutoTokenizer.from_pretrained(f"facebook/m2m100_1.2B")

model = MultiLingualTranslatorCT2fromHfHub(
    model_name_or_path="michaelfeil/ct2fast-m2m100_1.2B", device="cpu", compute_type="int8",
    tokenizer=tokenizer
)

'''
outputs = model.generate(
    ["In this scenario, we see three individuals - two women and one man. The first woman (W1) is walking on footpath while holding her mobile phone in hand with earphones plugged into it. She seems deep in thought or enjoying music from her device. Behind W1, there's another female pedestrian (W2), who appears slightly distracted as she looks around while strolling alongside a tree-lined pathway filled with greenery. Meanwhile, standing by himself near the edge of a cliff overlooking vast fields full of crops like wheat and corn, dressed formally but looking somewhat uncertain/anxious is the third individual, represented by the last emoji depicting him buttoning his suit jacket"],
    src_lang=["en"],
    tgt_lang=["zh"]
)
'''

example_sample = [
    ["What is the US currency?","en","zh"],
    ["What is the US currency?","en","ja"],
    ["美国的通货是什么？", "zh", "en"]
]

def demo_func(src_question, src_lang, tgt_lang):
    assert type(src_question) == type("")
    assert type(src_lang) == type("")
    assert type(tgt_lang) == type("")
    if "[SEP]" in src_question:
        src_question = list(filter(lambda xx: xx ,map(lambda x: x.strip() ,
        src_question.split("[SEP]"))))
    else:
        src_question = [src_question]
    '''
    tgt_question = trans_model.translate(
        src_question,
        source_lang=src_lang, target_lang = tgt_lang
    )
    '''
    tgt_question = model.generate(
    src_question,
    src_lang=[src_lang],
    tgt_lang=[tgt_lang]
    )
    assert type(tgt_question) == type([])
    tgt_question = "[SEP]".join(tgt_question)
    return {
        "Target Question": tgt_question
    }


demo = gr.Interface(
        fn=demo_func,
        inputs=[gr.Text(label = "Source Question"),
                gr.Text(label = "Source Language", value = "en"),
                gr.Text(label = "Target Language", value = "zh")
        ],
        outputs="json",
        title=f"cTranslate 🐱⚡️ demonstration",
        examples=example_sample if example_sample else None,
        cache_examples = False
    )

demo.launch(server_name=None, server_port=None)