Update app.py

app.py

@@ -1,37 +1,82 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# 加载稳定的中英互译模型
-translator_zh2en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en", device=-1)
-translator_en2zh = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh", device=-1)
-
-# 4种翻译场景
+# --------------------------
+# 1. 配置差异化场景（明确风格指令）
+# --------------------------
 SCENARIOS = {
-    "日常对话": "（口语化、简洁自然）",
-    "文档翻译": "（正式、准确，保留专业术语）",
-    "旅游场景": "（突出地点、时间等关键信息）",
-    "商务沟通": "（礼貌、专业的商务语境）"
+    "日常闲聊": "请翻译成非常口语化的表达，像朋友聊天一样，用轻松的语气：\n",
+    "正式文档": "请翻译成严谨的书面语，保留所有专业术语，语气正式：\n",
+    "旅游口语": "请翻译成突出地点/时间的旅游常用语，用词简单、语气轻松：\n",
+    "商务邮件": "请翻译成礼貌的商务用语，使用正式敬语，体现专业性：\n",
+    "网络热梗": "请翻译成当下流行的网络用语（可以用缩写/流行词）：\n",
+    "儿童话术": "请翻译成简单易懂的儿童口语，用词幼稚化、语气可爱：\n"
 }
 
+# --------------------------
+# 2. 加载支持指令的翻译模型（适配风格）
+# --------------------------
+# 选择支持多语言+指令的mbart模型（效果比opus-mt更适配风格）
+model_name = "facebook/mbart-large-50-many-to-many-mmt"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# --------------------------
+# 3. 翻译函数（强化场景指令）
+# --------------------------
 def translate(text, src_lang, tgt_lang, scenario):
     if not text.strip():
-        return "请输入内容"
-    prompt = f"{SCENARIOS[scenario]}\n{text}"
-    if src_lang == "中文" and tgt_lang == "英文":
-        return translator_zh2en(prompt)[0]["translation_text"]
-    else:
-        return translator_en2zh(prompt)[0]["translation_text"]
+        return "请输入要翻译的内容~"
+    
+    # 拼接“场景指令 + 待翻译文本”
+    prompt = f"{SCENARIOS[scenario]}{text}"
+    
+    # 设置模型的源语言/目标语言（mbart的语言编码）
+    lang_map = {
+        "中文": "zh_CN",
+        "英文": "en_XX"
+    }
+    src_code = lang_map[src_lang]
+    tgt_code = lang_map[tgt_lang]
+    tokenizer.src_lang = src_code  # 设定源语言
+    
+    # 模型推理
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id[tgt_code],  # 设定目标语言
+        max_new_tokens=512,
+        temperature=0.7  # 适当增加随机性，强化风格差异
+    )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 界面
-with gr.Blocks(title="多场景翻译助手") as demo:
-    gr.Markdown("# 多场景中英翻译工具")
+# --------------------------
+# 4. Gradio界面配置
+# --------------------------
+with gr.Blocks(title="多场景翻译工具") as demo:
+    gr.Markdown("# 📝 多风格翻译助手（差异更明显）")
     with gr.Row():
-        src = gr.Dropdown(["中文", "英文"], label="源语言", value="中文")
-        tgt = gr.Dropdown(["英文", "中文"], label="目标语言", value="英文")
-        scene = gr.Dropdown(list(SCENARIOS.keys()), label="翻译场景", value="日常对话")
-    input_box = gr.Textbox(label="输入", lines=4)
-    output_box = gr.Textbox(label="输出", lines=4)
-    gr.Button("翻译").click(translate, [input_box, src, tgt, scene], output_box)
+        with gr.Column(scale=2):
+            text_input = gr.Textbox(label="输入文本", lines=5, placeholder="在这里输入要翻译的内容...")
+            src_lang = gr.Radio(["中文", "英文"], label="源语言", value="中文")
+            tgt_lang = gr.Radio(["英文", "中文"], label="目标语言", value="英文")
+        with gr.Column(scale=1):
+            scenario = gr.Dropdown(
+                list(SCENARIOS.keys()),
+                label="翻译场景",
+                value="日常闲聊"
+            )
+            translate_btn = gr.Button("开始翻译", variant="primary")
+    text_output = gr.Textbox(label="翻译结果", lines=5)
+
+    # 绑定函数
+    translate_btn.click(
+        fn=translate,
+        inputs=[text_input, src_lang, tgt_lang, scenario],
+        outputs=text_output
+    )
 
+# 启动界面
 if __name__ == "__main__":
     demo.launch()