Update app.py

app.py

@@ -8,7 +8,6 @@ import time
 import numpy as np
 
 # 初始化模型
-@gr.cache
 def load_models():
     """加载三个不同的文本生成模型"""
     models = {}
@@ -16,33 +15,53 @@ def load_models():
     try:
         # 模型1: GPT-2 (轻量级)
         models['gpt2'] = {
-            'pipeline': pipeline("text-generation", model="gpt2", max_length=100),
+            'pipeline': pipeline("text-generation", model="gpt2", max_new_tokens=50),
             'name': 'GPT-2',
             'description': '经典的自回归语言模型，适合短文本生成'
         }
         
         # 模型2: DistilGPT-2 (更快速)
         models['distilgpt2'] = {
-            'pipeline': pipeline("text-generation", model="distilgpt2", max_length=100),
+            'pipeline': pipeline("text-generation", model="distilgpt2", max_new_tokens=50),
             'name': 'DistilGPT-2',
             'description': '轻量化的GPT-2，速度更快但质量略低'
         }
         
-        # 模型3: Microsoft DialoGPT (对话优化)
-        models['dialogpt'] = {
-            'pipeline': pipeline("text-generation", model="microsoft/DialoGPT-medium", max_length=100),
-            'name': 'DialoGPT-medium',
-            'description': '针对对话场景优化的生成模型'
+        # 模型3: OpenELM (苹果开源模型)
+        models['openelm'] = {
+            'pipeline': pipeline("text-generation", model="apple/OpenELM-270M", max_new_tokens=50, trust_remote_code=True),
+            'name': 'OpenELM-270M',
+            'description': '苹果开源的轻量级语言模型'
         }
         
     except Exception as e:
         print(f"模型加载错误: {e}")
-        # 备用方案：使用更简单的模型
-        models['gpt2'] = {
-            'pipeline': pipeline("text-generation", model="gpt2", max_length=50),
-            'name': 'GPT-2',
-            'description': '经典的自回归语言模型'
-        }
+        # 备用方案：只使用最基础的模型
+        try:
+            models['gpt2'] = {
+                'pipeline': pipeline("text-generation", model="gpt2", max_new_tokens=30),
+                'name': 'GPT-2',
+                'description': '经典的自回归语言模型'
+            }
+            models['distilgpt2'] = {
+                'pipeline': pipeline("text-generation", model="distilgpt2", max_new_tokens=30),
+                'name': 'DistilGPT-2',
+                'description': '轻量化版本'
+            }
+            # 第三个模型用简单的替代
+            models['openelm'] = {
+                'pipeline': pipeline("text-generation", model="gpt2", max_new_tokens=20),
+                'name': 'GPT-2-Variant',
+                'description': '备用模型配置'
+            }
+        except Exception as e2:
+            print(f"备用模型加载也失败: {e2}")
+            # 最终备用：至少确保有一个模型可用
+            models['gpt2'] = {
+                'pipeline': None,
+                'name': 'GPT-2',
+                'description': '模型加载失败'
+            }
     
     return models
 
@@ -63,35 +82,41 @@ GRACE_DATA = {
         'Artistry': 6.8,
         'Efficiency': 9.2
     },
-    'DialoGPT-medium': {
-        'Generalization': 7.0,
-        'Relevance': 8.8,
-        'Artistry': 8.0,
-        'Efficiency': 7.5
+    'OpenELM-270M': {
+        'Generalization': 6.5,
+        'Relevance': 7.0,
+        'Artistry': 6.5,
+        'Efficiency': 8.8
     }
 }
 
-def generate_text_with_model(model_key, prompt, max_length=100):
+def generate_text_with_model(model_key, prompt, max_length=50):
     """使用指定模型生成文本"""
     try:
         start_time = time.time()
         
-        if model_key not in MODELS:
-            return "模型未找到", 0
+        if model_key not in MODELS or MODELS[model_key]['pipeline'] is None:
+            return "模型未找到或加载失败", 0
             
         result = MODELS[model_key]['pipeline'](
             prompt, 
-            max_length=max_length,
+            max_new_tokens=min(max_length, 50),
             num_return_sequences=1,
             temperature=0.7,
             do_sample=True,
-            pad_token_id=50256
+            pad_token_id=50256,
+            truncation=True,
+            return_full_text=False
         )
         
         end_time = time.time()
         generation_time = end_time - start_time
         
-        generated_text = result[0]['generated_text']
+        if result and len(result) > 0:
+            generated_text = prompt + result[0]['generated_text']
+        else:
+            generated_text = "生成失败"
+            
         return generated_text, generation_time
         
     except Exception as e:
@@ -161,7 +186,7 @@ def arena_interface(prompt, max_length):
     # 格式化输出
     output1 = f"**{MODELS['gpt2']['name']}** (生成时间: {times.get('gpt2', 0):.2f}s)\n\n{results.get('gpt2', '生成失败')}"
     output2 = f"**{MODELS['distilgpt2']['name']}** (生成时间: {times.get('distilgpt2', 0):.2f}s)\n\n{results.get('distilgpt2', '生成失败')}"
-    output3 = f"**{MODELS['dialogpt']['name']}** (生成时间: {times.get('dialogpt', 0):.2f}s)\n\n{results.get('dialogpt', '生成失败')}"
+    output3 = f"**{MODELS['openelm']['name']}** (生成时间: {times.get('openelm', 0):.2f}s)\n\n{results.get('openelm', '生成失败')}"
     
     # 生成对比分析
     analysis = f"""
@@ -170,13 +195,13 @@ def arena_interface(prompt, max_length):
 ### 速度对比
 - GPT-2: {times.get('gpt2', 0):.2f}秒
 - DistilGPT-2: {times.get('distilgpt2', 0):.2f}秒  
-- DialoGPT: {times.get('dialogpt', 0):.2f}秒
+- OpenELM: {times.get('openelm', 0):.2f}秒
 
 ### 质量评估
 根据GRACE框架，不同模型在各维度的表现存在差异：
-- **效率性**: DistilGPT-2表现最佳
-- **相关性**: DialoGPT在对话场景中表现突出
+- **效率性**: DistilGPT-2和OpenELM表现优异
 - **泛化性**: GPT-2具有最强的通用性
+- **相关性**: 各模型在相关性上表现相近
 """
     
     return output1, output2, output3, analysis
@@ -239,7 +264,7 @@ def create_app():
                 with gr.Row():
                     model1_output = gr.Markdown(label="GPT-2 输出")
                     model2_output = gr.Markdown(label="DistilGPT-2 输出")
-                    model3_output = gr.Markdown(label="DialoGPT 输出")
+                    model3_output = gr.Markdown(label="OpenELM 输出")
                 
                 analysis_output = gr.Markdown(label="对比分析")