Add transjob v1.0

app.py

@@ -1,17 +1,41 @@
 from transformers import MT5ForConditionalGeneration, AutoTokenizer
-import gradio as grad
+import gradio as gr
 
-mdl = MT5ForConditionalGeneration.from_pretrained("google/mt5-small")
-tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
+trans_mdl = MT5ForConditionalGeneration.from_pretrained("K024/mt5-zh-ja-en-trimmed")
+trans_tokenizer = AutoTokenizer.from_pretrained("K024/mt5-zh-ja-en-trimmed")
 
-def translation_CN2EN(text):
-    inp =  "translate Chinese to English: " + text
-    enc = tokenizer(inp, return_tensors="pt").input_ids
-    tokens = mdl(**enc)
-    response = tokenizer.decode(tokens[0], skip_special_tokens=True)
+
+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
-para=grad.Textbox(lines=1, label="Chinese Text", placeholder="Text in Chinese")
-out=grad.Textbox(lines=1, label="English Translation")
-grad.Interface(translation_CN2EN, inputs=para, outputs=out).launch()
+    
+
+
+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()