import gradio as gr

from transformers import pipeline


model_list = [
    "facebook/m2m100_418M", 
    "Helsinki-NLP/opus-mt-en-zh", 
    "liam168/trans-opus-mt-en-zh"
    ]


translation_pipelines = {}
# init
for model_name in model_list:
    translation_pipelines[model_name] = pipeline("translation_en_to_zh", model=model_name)

def translate(text, model_name):
    pipe = translation_pipelines[model_name]
    return pipe(text)[0]["translation_text"] if text.strip() else ''

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            englishtxt_model_input = [
                gr.Textbox(label="English text"), 
                gr.Dropdown(model_list, value="facebook/m2m100_418M", label="Model")
            ]
            translate_btn = gr.Button(value="Translate")
        with gr.Column():
            chinese_output = gr.Textbox(label="Chinese Text")

    translate_btn.click(translate, inputs=englishtxt_model_input, outputs=chinese_output, api_name="translate-to-chinese")
    examples = gr.Examples(examples=["I went to the supermarket yesterday.", "Helen is a good swimmer."],
                           inputs=[englishtxt_model_input[0]])

demo.launch()