from easynmt import EasyNMT
import os

trans_model = EasyNMT("m2m_100_1.2B")

import gradio as gr

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
    )
    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"Translate 🐱 demonstration",
        examples=example_sample if example_sample else None,
        cache_examples = False
    )

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