from transformers import MT5ForConditionalGeneration, AutoTokenizer
import gradio as gr

trans_mdl = MT5ForConditionalGeneration.from_pretrained("K024/mt5-zh-ja-en-trimmed")
trans_tokenizer = AutoTokenizer.from_pretrained("K024/mt5-zh-ja-en-trimmed")


def translation_job(data):
    job = data[0]
    text = data[1]


    # 设置翻译任务和提示语的映射
    job_key = ["中译日", "中译英", "日译中", "英译中", "日译英", "英译日"]
    job_value = ["zh2ja:", "zh2en:", "ja2zh:", "en2zh:", "ja2en:", "en2ja:"]
    job_map = dict(zip(job_key, job_value))

    input = job_map[job] + text
    enc = trans_tokenizer(input, return_tensor="pt")
    tokens = trans_mdl.generate(**enc)
    response = trans_tokenizer.batch_decode(tokens)
    return response
    


with gr.Blocks() as app:
    # 中英日三语翻译任务
    with gr.Tab("中英日三语翻译"):
        job_name = gr.Dropdown(
            ["中译日", "中译英", "日译中", "英译中", "日译英", "英译日"],
            label = "翻译任务选择",
            info = "单选"
        )
        source_text = gr.Textbox(lines=1, label="翻译文本", placeholder="请输入要翻译的文本")
        trans_result = gr.Textbox(lines=1, label="翻译结果", placeholder="翻译结果")
        trans_btn = gr.Button("翻译")

    trans_btn.click(translation_job, inputs=[job_name, source_text], outputs=trans_result)
app.launch()