Update app.py

app.py

@@ -9,20 +9,27 @@ import torch
 import json
 import re
 
-# 选择两个翻译模型
+# 选择两个中文到英文的翻译模型
 MODEL_CONFIGS = {
-    "English-to-Chinese": {
-        "model_name": "Helsinki-NLP/opus-mt-en-zh",
-        "description": "英文到中文的机器翻译模型 (Helsinki-NLP OPUS-MT)",
+    "Chinese-to-English (Opus-MT)": {
+        "model_name": "Helsinki-NLP/opus-mt-zh-en",
+        "description": "中文到英文的机器翻译模型 (Helsinki-NLP OPUS-MT)",
         "max_length": 200, # 翻译输出的最大长度
         "color": "#FF6B6B"
     },
-    "Chinese-to-English": {
-        "model_name": "Helsinki-NLP/opus-mt-zh-en",
-        "description": "中文到英文的机器翻译模型 (Helsinki-NLP OPUS-MT)",
+    "Chinese-to-English (M4-Small)": {
+        "model_name": "HuggingFaceM4/m4-small-en-zh", # 这是一个多语言模型，支持zh-en
+        "description": "中文到英文的机器翻译模型 (HuggingFaceM4 M4-Small)",
         "max_length": 200, # 翻译输出的最大长度
         "color": "#4ECDC4"
     }
+    # 如果需要第三个模型，可以取消注释下面这个，或替换成您想要的
+    # "Chinese-to-English (Another Model)": {
+    #     "model_name": "facebook/mbart-large-50-one-to-many-mmt", # 另一个多语言模型，需要指定 src_lang/tgt_lang
+    #     "description": "中文到英文的机器翻译模型 (Facebook mBART-Large-50)",
+    #     "max_length": 200,
+    #     "color": "#45B7D1"
+    # }
 }
 
 class TranslationComparator:
@@ -37,13 +44,36 @@ class TranslationComparator:
             try:
                 print(f"加载 {model_key} ({config['model_name']})...")
                 # 对于翻译任务，使用 "translation" pipeline
-                # Gradio-spaces 上的 free tier 可能会导致内存不足，所以 device=-1 (CPU) 更稳妥
+                # 注意：某些多语言模型（如 m4-small）可能需要显式指定源语言和目标语言
+                # 对于 Helsinki-NLP/opus-mt-zh-en，pipeline会自动处理
+                # 对于 HuggingFaceM4/m4-small-en-zh，虽然名字是en-zh，但它内部支持zh-en。
+                # 如果遇到问题，可能需要更复杂的tokenizer/model加载方式而非pipeline
+                if "opus-mt-zh-en" in config["model_name"]:
+                    task = "translation_zh_to_en" # 更明确的翻译任务
+                elif "m4-small" in config["model_name"]:
+                    # m4-small是一个多语言模型，需要提供源语言和目标语言。
+                    # pipeline("translation") 不直接支持 src_lang/tgt_lang 参数
+                    # 需要手动加载 AutoModelForSeq2SeqLM 和 AutoTokenizer
+                    print(f"特别加载 {model_key} 及其Tokenizer...")
+                    tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
+                    model = AutoModelForSeq2SeqLM.from_pretrained(config["model_name"])
+                    # 将其包装成一个简单的可调用对象，模拟pipeline的行为
+                    self.models[model_key] = {
+                        "tokenizer": tokenizer,
+                        "model": model,
+                        "pipeline_type": "custom_translation"
+                    }
+                    print(f"✓ {model_key} 加载成功 (自定义翻译模式)")
+                    continue # 跳过pipeline加载
+                else: # 默认翻译任务
+                    task = "translation"
+
                 self.models[model_key] = pipeline(
-                    "translation",
+                    task,
                     model=config["model_name"],
                     tokenizer=config["model_name"],
                     device=-1, # 使用CPU
-                    torch_dtype=torch.float32 # 保持一致，或根据模型精度调整
+                    torch_dtype=torch.float32
                 )
                 print(f"✓ {model_key} 加载成功")
             except Exception as e:
@@ -52,7 +82,8 @@ class TranslationComparator:
 
     def translate_text(self, model_key, text_to_translate, max_length=200):
         """使用指定模型进行翻译"""
-        if self.models[model_key] is None:
+        model_entry = self.models.get(model_key)
+        if model_entry is None:
             return {
                 "translated_text": f"[Model {model_key} not loaded correctly, this is a simulated translation]",
                 "inference_time": 0.5,
@@ -65,18 +96,38 @@ class TranslationComparator:
 
         try:
             start_time = time.time()
-            
-            # 翻译文本
-            # pipeline("translation") 的返回格式是 [{"translation_text": "..."}]
-            result = self.models[model_key](
-                text_to_translate,
-                max_length=max_length
-            )
+
+            if isinstance(model_entry, dict) and model_entry.get("pipeline_type") == "custom_translation":
+                # 对于需要自定义处理的模型 (如 HuggingFaceM4/m4-small-en-zh)
+                tokenizer = model_entry["tokenizer"]
+                model = model_entry["model"]
+
+                # 对于 m4-small，需要手动设置源语言和目标语言
+                # 假设输入是中文
+                input_ids = tokenizer(text_to_translate, return_tensors="pt", truncation=True, max_length=512).input_ids
+                
+                # 设置生成参数，特别是强制生成目标语言的 token (en_XX)
+                # 对于 m4-small 而言，`en_XX` 是英文的目标语言token
+                # 请注意：这可能需要根据具体的m4模型进行微调，因为它可能没有直接的force_bos_token_id
+                # 一个更通用的方法是手动构建decoder_input_ids
+                
+                # 尝试一个通用的生成方式，让模型自己识别语言
+                # 对于翻译任务，transformers pipeline已经封装了大部分复杂性
+                # 如果手动调用generate，需要确保输入格式和语言ID正确
+                
+                # 简单的直接生成（可能不带force_bos_token_id）
+                generated_ids = model.generate(input_ids, max_new_tokens=max_length)
+                translated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+
+            else: # 使用 pipeline
+                result = model_entry(
+                    text_to_translate,
+                    max_length=max_length
+                )
+                translated_text = result[0]['translation_text']
             
             end_time = time.time()
             
-            translated_text = result[0]['translation_text']
-            
             return {
                 "translated_text": translated_text,
                 "inference_time": round(end_time - start_time, 3),
@@ -98,84 +149,59 @@ class TranslationComparator:
 # 初始化比较器
 comparator = TranslationComparator()
 
-def run_translation_comparison(en_prompt, zh_prompt, max_length):
-    """运行所有模型的翻译对比"""
+def run_translation_comparison(zh_prompt, max_length):
+    """运行所有中文到英文模型的翻译对比"""
     
-    results = {}
+    if not zh_prompt.strip():
+        # 返回与模型数量相匹配的错误消息
+        return tuple([gr.Code.update(value=json.dumps({"错误信息": "请输入中文文本进行翻译"}, ensure_ascii=False)) for _ in MODEL_CONFIGS])
 
-    # 英文到中文翻译
-    if "English-to-Chinese" in MODEL_CONFIGS and en_prompt.strip():
-        result_en_zh = comparator.translate_text(
-            "English-to-Chinese",
-            en_prompt,
-            max_length=int(max_length)
-        )
-        results["English-to-Chinese"] = result_en_zh
-    else:
-        results["English-to-Chinese"] = {"error": "请输入英文文本进行翻译"} if not en_prompt.strip() else {}
+    results = {}
+    outputs_list = []
 
-    # 中文到英文翻译
-    if "Chinese-to-English" in MODEL_CONFIGS and zh_prompt.strip():
-        result_zh_en = comparator.translate_text(
-            "Chinese-to-English",
+    for model_key in MODEL_CONFIGS.keys():
+        result = comparator.translate_text(
+            model_key,
             zh_prompt,
             max_length=int(max_length)
         )
-        results["Chinese-to-English"] = result_zh_en
-    else:
-        results["Chinese-to-English"] = {"error": "请输入中文文本进行翻译"} if not zh_prompt.strip() else {}
-
-    # 格式化输出
-    def format_result(result):
+        results[model_key] = result
+    
+        # 格式化输出
         if "error" in result:
-            return json.dumps({"错误信息": result["error"]}, indent=2, ensure_ascii=False)
-        
-        formatted = {
-            "翻译文本": result["translated_text"],
-            "推断时间": f"{result['inference_time']}s",
-            "翻译Token数": result["output_length"],
-            "翻译速度": f"{result['output_length']/max(result['inference_time'], 0.001):.1f} tokens/s"
-        }
-        return json.dumps(formatted, indent=2, ensure_ascii=False)
-
-    en_zh_output = format_result(results.get("English-to-Chinese", {}))
-    zh_en_output = format_result(results.get("Chinese-to-English", {}))
+            outputs_list.append(json.dumps({"错误信息": result["error"]}, indent=2, ensure_ascii=False))
+        else:
+            formatted = {
+                "翻译文本": result["translated_text"],
+                "推断时间": f"{result['inference_time']}s",
+                "翻译Token数": result["output_length"],
+                "翻译速度": f"{result['output_length']/max(result['inference_time'], 0.001):.1f} tokens/s"
+            }
+            outputs_list.append(json.dumps(formatted, indent=2, ensure_ascii=False))
     
-    # 假设有第三个模型（如果 MODEL_CONFIGS 扩展了）
-    # 如果没有，这将是空字符串
-    third_model_key = list(MODEL_CONFIGS.keys())[2] if len(MODEL_CONFIGS) > 2 else None
-    third_output = format_result(results.get(third_model_key, {})) if third_model_key else ""
-
-
-    # 动态调整返回的输出数量
-    if len(MODEL_CONFIGS) == 2:
-        return en_zh_output, zh_en_output
-    else: # 假设有3个模型
-        return en_zh_output, zh_en_output, third_output
+    return tuple(outputs_list)
 
 
 def calculate_grace_scores_for_translation():
     """为翻译任务计算GRACE评估分数"""
+    # 模拟中文到英文翻译模型的GRACE分数
     grace_data = {
-        "English-to-Chinese": {
-            "Generalization": 8.0, # 处理不同领域英翻中能力
-            "Relevance": 8.5,     # 翻译内容与原文语义相关性
-            "Accuracy": 8.2,      # 翻译精确性
-            "Consistency": 8.0,   # 翻译稳定性
-            "Efficiency": 7.5     # 推理效率
-        },
-        "Chinese-to-English": {
+        "Chinese-to-English (Opus-MT)": {
             "Generalization": 7.8, # 处理不同领域中翻英能力
             "Relevance": 8.3,     # 翻译内容与原文语义相关性
             "Accuracy": 8.0,      # 翻译精确性
             "Consistency": 7.9,   # 翻译稳定性
             "Efficiency": 7.5     # 推理效率
+        },
+        "Chinese-to-English (M4-Small)": {
+            "Generalization": 7.0, # 多语言模型可能在特定语对上略逊色于专用模型
+            "Relevance": 7.5,
+            "Accuracy": 7.2,
+            "Consistency": 7.0,
+            "Efficiency": 8.5     # 通常小模型效率更高
         }
+        # 如果有第三个模型，在这里添加其分数
     }
-    # 如果有第三个模型，可以添加其分数
-    # "Another-Translation-Model": {
-    #     "Generalization": ..., "Relevance": ..., "Accuracy": ..., "Consistency": ..., "Efficiency": ...
-    # }
     return grace_data
 
 
@@ -210,7 +236,7 @@ def create_translation_radar_chart():
         ),
         showlegend=True,
         title={
-            'text': "GRACE框架：翻译模型评估",
+            'text': "GRACE框架：中文到英文翻译模型评估",
             'x': 0.5,
             'font': {'size': 16}
         },
@@ -239,7 +265,7 @@ def create_performance_bar_chart():
         ))
 
     fig.update_layout(
-        title='GRACE框架详细对比 - 翻译',
+        title='GRACE框架详细对比 - 中文到英文翻译',
         xaxis_title='模型',
         yaxis_title='分数 (0-10)',
         barmode='group',
@@ -252,9 +278,16 @@ def create_model_info_table():
     """创建模型信息对比表"""
     model_info = []
     for model_key, config in MODEL_CONFIGS.items():
-        # 模拟参数信息 (Helsinki-NLP OPUS-MT 模型通常较小)
-        params = "~3亿" if "en-zh" in config["model_name"] else "~3亿"
-        size = "~1.2GB" if "en-zh" in config["model_name"] else "~1.2GB"
+        # 模拟参数信息
+        if "opus-mt-zh-en" in config["model_name"]:
+            params = "~3亿"
+            size = "~1.2GB"
+        elif "m4-small" in config["model_name"]:
+            params = "~4亿" # m4-small 实际参数量可能更大
+            size = "~1.5GB"
+        else: # 默认值
+            params = "未知"
+            size = "未知"
         
         model_info.append({
             "模型": model_key,
@@ -283,84 +316,64 @@ def create_summary_scores_table():
     df = pd.DataFrame(summary_data)
     return df
 
-# 预设的示例提示（英文和中文）
-EXAMPLE_EN_PROMPTS = [
-    "Hello, how are you today?",
-    "The quick brown fox jumps over the lazy dog.",
-    "Artificial intelligence is transforming many industries."
-]
-
+# 预设的示例中文提示
 EXAMPLE_ZH_PROMPTS = [
     "你好，今天过得怎么样？",
     "敏捷的棕色狐狸跳过懒惰的狗。",
-    "人工智能正在改变许多行业。"
+    "人工智能正在改变许多行业。",
+    "今天天气真好，我们去公园散步吧。"
 ]
 
 def create_app():
-    with gr.Blocks(title="翻译模型对比", theme=gr.themes.Soft()) as app:
-        gr.Markdown("# 🌐 翻译模型对比竞技场")
-        gr.Markdown("### 使用GRACE框架对比不同翻译模型在翻译任务中的表现")
+    with gr.Blocks(title="中文到英文翻译模型对比", theme=gr.themes.Soft()) as app:
+        gr.Markdown("# 🌐 中文到英文翻译模型对比竞技场")
+        gr.Markdown("### 使用GRACE框架对比不同中文到英文翻译模型在翻译任务中的表现")
         
         with gr.Tabs():
             # Arena选项卡
             with gr.TabItem("️ 翻译竞技场"):
                 gr.Markdown("## 翻译竞技场")
-                gr.Markdown("请在下方输入需要翻译的文本（英文或中文），查看不同模型的翻译效果。")
+                gr.Markdown("请在下方输入需要翻译的**中文**文本，查看不同模型翻译成**英文**的效果。")
                 
                 with gr.Row():
-                    with gr.Column(scale=1):
-                        input_en_prompt = gr.Textbox(
-                            label="输入英文文本",
-                            placeholder="Enter your English text here...",
-                            lines=3,
-                            value=EXAMPLE_EN_PROMPTS[0]
-                        )
-                        # 预设英文示例按钮
-                        with gr.Row():
-                            for i, example in enumerate(EXAMPLE_EN_PROMPTS):
-                                gr.Button(f"英文示例 {i+1}", size="sm").click(
-                                    fn=lambda x=example: x,
-                                    outputs=[input_en_prompt]
-                                )
-                    
-                    with gr.Column(scale=1):
+                    with gr.Column(scale=2): # 增加输入框的比例
                         input_zh_prompt = gr.Textbox(
                             label="输入中文文本",
                             placeholder="在此输入您的中文文本...",
-                            lines=3,
+                            lines=4, # 增加行数
                             value=EXAMPLE_ZH_PROMPTS[0]
                         )
                         # 预设中文示例按钮
                         with gr.Row():
                             for i, example in enumerate(EXAMPLE_ZH_PROMPTS):
-                                gr.Button(f"中文示例 {i+1}", size="sm").click(
+                                gr.Button(f"示例 {i+1}", size="sm").click(
                                     fn=lambda x=example: x,
                                     outputs=[input_zh_prompt]
                                 )
-
-                with gr.Column(scale=1):
-                    max_length = gr.Slider(
-                        minimum=50,
-                        maximum=500,
-                        value=200,
-                        step=10,
-                        label="最大输出Token数"
-                    )
                     
-                    submit_btn = gr.Button(" 开始翻译", variant="primary", size="lg")
+                    with gr.Column(scale=1): # 调整参数控制列的比例
+                        max_length = gr.Slider(
+                            minimum=50,
+                            maximum=500,
+                            value=200,
+                            step=10,
+                            label="最大输出Token数"
+                        )
+                        
+                        submit_btn = gr.Button(" 开始翻译", variant="primary", size="lg")
                 
                 # 动态创建输出框
                 output_boxes = []
                 for model_key, config in MODEL_CONFIGS.items():
                     output_boxes.append(gr.Code(
-                        label=f"{model_key} ({config['description'].split('(')[0].strip()})",
+                        label=f"{model_key} 翻译结果", # 明确翻译方向
                         language="json",
                         value="点击“开始翻译”查看结果"
                     ))
 
                 submit_btn.click(
                     fn=run_translation_comparison,
-                    inputs=[input_en_prompt, input_zh_prompt, max_length],
+                    inputs=[input_zh_prompt, max_length],
                     outputs=output_boxes
                 )