Update app.py

app.py

@@ -4,7 +4,8 @@ import plotly.graph_objects as go
 import plotly.express as px
 import time
 import numpy as np
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+# 导入 AutoTokenizer 和 AutoModelForSeq2SeqLM
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM 
 import torch
 import json
 import re
@@ -23,6 +24,13 @@ MODEL_CONFIGS = {
         "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:
@@ -46,8 +54,8 @@ class TranslationComparator:
                         tokenizer=config["model_name"],
                         src_lang="zh_CN",  # 源语言为中文
                         tgt_lang="en_US",  # 目标语言为英文
-                        device=-1, # 使用CPU
-                        torch_dtype=torch.float32
+                        device=-1, # 使用CPU，避免GPU内存不足问题
+                        torch_dtype=torch.float32 # 保持一致，或根据模型精度调整
                     )
                 else: # 对于Helsinki-NLP/opus-mt-zh-en等
                     self.models[model_key] = pipeline(
@@ -78,38 +86,18 @@ class TranslationComparator:
 
         try:
             start_time = time.time()
-
-            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']
+            
+            # 翻译文本
+            # pipeline("translation") 的返回格式是 [{"translation_text": "..."}]
+            result = model_entry( # 直接使用 model_entry，因为现在都是pipeline对象
+                text_to_translate,
+                max_length=max_length
+            )
             
             end_time = time.time()
             
+            translated_text = result[0]['translation_text']
+            
             return {
                 "translated_text": translated_text,
                 "inference_time": round(end_time - start_time, 3),
@@ -175,14 +163,13 @@ def calculate_grace_scores_for_translation():
             "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     # 通常小模型效率更高
+        "Chinese-to-English (mBART-Large-50)": { # **这里已修改！**
+            "Generalization": 8.5, # 更大型多语言模型，泛化性通常更强
+            "Relevance": 8.8,
+            "Accuracy": 8.6,
+            "Consistency": 8.5,
+            "Efficiency": 6.0     # 模型较大，效率可能略低
         }
-        # 如果有第三个模型，在这里添加其分数
     }
     return grace_data
 
@@ -196,7 +183,9 @@ def create_translation_radar_chart():
 
     for i, (model_name, scores) in enumerate(grace_scores.items()):
         values = [scores[cat] for cat in categories]
-        color = MODEL_CONFIGS[model_name]["color"]
+        # **这里使用 MODEL_CONFIGS[model_name]["color"] 依赖于 MODEL_CONFIGS 和 grace_scores 的键名一致**
+        # 这是导致之前 KeyError 的地方，现在应该已修复，因为 calculate_grace_scores_for_translation 的键名已更新
+        color = MODEL_CONFIGS[model_name]["color"] 
 
         fig.add_trace(go.Scatterpolar(
             r=values,
@@ -264,9 +253,9 @@ def create_model_info_table():
         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"
+        elif "mbart-large-50" in config["model_name"]: # 修改为mBART的参数
+            params = "~6.1亿" # mBART-Large-50 的实际参数量
+            size = "~2.4GB" # mBART-Large-50 的实际模型大小
         else: # 默认值
             params = "未知"
             size = "未知"