Update app.py

app.py

@@ -1,12 +1,18 @@
 import os
 import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, WhisperProcessor, WhisperForConditionalGeneration
+from typing import List, Tuple  # 新增：导入类型
 
-# 从环境变量获取Hugging Face Token
+# 方案 A：使用自定义环境变量名 "language"
 hf_token = os.environ.get("language")
 if not hf_token:
-    raise EnvironmentError("未找到HUGGINGFACE_HUB_TOKEN环境变量，请在Space设置中添加")
+    raise EnvironmentError("未找到名为 'language' 的环境变量，请在Space设置中添加")
+
+# 方案 B：改用规范的 "HUGGINGFACE_HUB_TOKEN"（需同步修改Space环境变量）
+# hf_token = os.environ.get("HUGGINGFACE_HUB_TOKEN")
+# if not hf_token:
+#     raise EnvironmentError("未找到HUGGINGFACE_HUB_TOKEN环境变量，请在Space设置中添加")
 
 # 模型配置 - 使用公开模型
 MODELS = {
@@ -35,38 +41,29 @@ def load_model(model_name):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     return model.to(device), tokenizer, device
 
-# 其余代码（界面构建和交互逻辑）保持不变...
-
 # 初始化模型
 loaded_models = {}
 for model_name in MODELS:
     loaded_models[model_name] = load_model(model_name)
 
-# 构建对话提示词（针对不同模型可能需要不同格式）
-def build_prompt(message, history, system_prompt, model_name):
-    # Zephyr/Mistral等模型使用简单格式
-    if "Zephyr" in model_name or "Mistral" in model_name:
+# 构建对话提示词
+def build_prompt(message: str, history: List[Tuple[str, str]], system_prompt: str, model_name: str) -> str:
+    if "Zephyr" in model_name:
         prompt = f"系统提示: {system_prompt}\n"
         for user_msg, assistant_msg in history:
             prompt += f"用户: {user_msg}\n助手: {assistant_msg}\n"
         prompt += f"用户: {message}\n助手:"
-        return prompt
-    
-    # Falcon模型使用更简洁的格式
     elif "Falcon" in model_name:
         prompt = f"### System:\n{system_prompt}\n\n"
         for user_msg, assistant_msg in history:
             prompt += f"### User:\n{user_msg}\n\n### Assistant:\n{assistant_msg}\n\n"
         prompt += f"### User:\n{message}\n\n### Assistant:"
-        return prompt
-    
-    # 默认为通用格式
     else:
         prompt = f"[System] {system_prompt}\n"
         for user_msg, assistant_msg in history:
             prompt += f"[User] {user_msg}\n[Assistant] {assistant_msg}\n"
         prompt += f"[User] {message}\n[Assistant]"
-        return prompt
+    return prompt
 
 # 模型推理函数
 def generate_response(
@@ -78,16 +75,11 @@ def generate_response(
     temperature: float,
     top_p: float,
     top_k: int
-):
+) -> str:
     model, tokenizer, device = loaded_models[model_name]
-    
-    # 构建提示词
     full_prompt = build_prompt(message, history, system_prompt, model_name)
-    
-    # 编码输入
     inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
     
-    # 生成参数
     generate_kwargs = {
         "max_new_tokens": max_new_tokens,
         "temperature": temperature,
@@ -98,20 +90,10 @@ def generate_response(
         "pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id
     }
     
-    # 生成响应
     with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            **generate_kwargs
-        )
-    
-    # 解码输出
+        output = model.generate(**inputs, **generate_kwargs)
     response = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    # 提取模型生成的部分
-    response = response[len(full_prompt):].strip()
-    
-    return response
+    return response[len(full_prompt):].strip()
 
 # 处理用户输入
 def process_chat(
@@ -123,11 +105,8 @@ def process_chat(
     temperature: float,
     top_p: float,
     top_k: int
-):
-    response = generate_response(
-        message, history, system_prompt, model_name,
-        max_new_tokens, temperature, top_p, top_k
-    )
+) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
+    response = generate_response(message, history, system_prompt, model_name, max_new_tokens, temperature, top_p, top_k)
     history.append((message, response))
     return history, history
 
@@ -135,14 +114,14 @@ def process_chat(
 asr = None
 if torch.cuda.is_available() or torch.backends.mps.is_available():
     try:
-        from transformers import WhisperProcessor, WhisperForConditionalGeneration
         processor = WhisperProcessor.from_pretrained("openai/whisper-base")
         asr_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base").to("cuda" if torch.cuda.is_available() else "cpu")
         asr = {"processor": processor, "model": asr_model}
-    except:
+    except Exception as e:
+        print(f"语音模型加载失败: {e}")
         asr = None
 
-def transcribe(audio):
+def transcribe(audio) -> str:
     if asr is None:
         return "语音识别模型未加载"
     processor, model = asr["processor"], asr["model"]
@@ -189,8 +168,7 @@ with gr.Blocks(title="无权限语言模型对话助手") as demo:
     # 发送消息
     send_btn.click(
         fn=process_chat,
-        inputs=[message_input, chat_history, system_prompt, model_choice,
-                max_new_tokens, temperature, top_p, top_k],
+        inputs=[message_input, chat_history, system_prompt, model_choice, max_new_tokens, temperature, top_p, top_k],
         outputs=[chat_history, chat_history]
     )