from transformers import MT5ForConditionalGeneration, AutoTokenizer, Text2TextGenerationPipeline
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")
trans_pipe = Text2TextGenerationPipeline(model=trans_mdl, tokenizer=trans_tokenizer)

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

    input = job_map[job] + text
    print(input)
    response = trans_tokenizer.batch_decode(input, max_length=100, num_beams=4)
    return response[0]['generated_text']
    


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="翻译结果")
        trans_btn = gr.Button("翻译")

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