Update app.py

app.py

@@ -1,3 +1,7 @@
+# ===================================================================
+#         优化版 app.py (为 Gradio 5.x 优化)
+# ===================================================================
+
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -5,46 +9,69 @@ from threading import Thread
 import logging
 import time
 import json
+import os
 
-# ===================================================================
-#         最终版 app.py (适配 Gradio 4.x+ 的推荐模式)
-# ===================================================================
+# --- 1. 配置与初始化 ---
 
-# 1. 配置详细的日志记录
+# 日志记录配置
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-MODEL_ID = "badanwang/teacher_basic_qwen3-0.6b"
+# 从环境变量或默认值加载模型ID，增加灵活性
+MODEL_ID = os.getenv("MODEL_ID", "badanwang/teacher_basic_qwen3-0.6b")
 
 logger.info("===== Application Startup =====")
 logger.info(f"正在加载模型和分词器: {MODEL_ID}")
+
+# 异常处理以优雅地处理模型加载失败
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    # 推荐使用 trust_remote_code=True 以确保所有模型组件正确加载
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        torch_dtype=torch.float32, # 在 CPU 上使用 float32 以获得最佳稳定性和兼容性
-        device_map="auto"
+        torch_dtype="auto",  # 推荐使用 "auto" 以获得最佳性能和兼容性
+        device_map="auto",
+        trust_remote_code=True
     )
     logger.info("模型和分词器加载成功！")
 except Exception as e:
     logger.error(f"加载模型时发生致命错误: {e}", exc_info=True)
+    # 在无法加载模型时，创建一个占位符函数，以便Gradio界面仍能启动并显示错误
     model, tokenizer = None, None
+    def model_load_error_placeholder(*args, **kwargs):
+        raise gr.Error(f"致命错误：无法加载模型 '{MODEL_ID}'。请检查后台日志以获取详细信息。")
 
-# --- 2. 核心推理函数 (已根据 Gradio 新模式重构) ---
-def predict(messages: list):
+# --- 2. 核心推理函数 ---
+
+# Gradio 5.x 的 ChatInterface `fn` 函数接收两个参数: message 和 history
+def predict(message: str, history: list[list[str]]):
     """
-    接收一个包含完整对话历史的 OpenAI 格式列表，返回模型的流式响应。
+    核心推理函数，接收用户输入和聊天历史，并以流式方式返回模型输出。
+
+    Args:
+        message (str): 用户的最新输入。
+        history (list[list[str]]): 聊天历史，格式为 [[user_msg_1, bot_msg_1], [user_msg_2, bot_msg_2], ...]。
     """
     start_time = time.time()
     logger.info("\n--- [START] 进入 predict 函数 ---")
-    
-    # 使用 json.dumps 美化输出，方便阅读
-    logger.info(f"[INPUT] 收到的 messages 列表:\n{json.dumps(messages, indent=2, ensure_ascii=False)}")
+    logger.info(f"[INPUT] Message: {message}")
+    logger.info(f"[INPUT] History:\n{json.dumps(history, indent=2, ensure_ascii=False)}")
 
+    # 如果模型加载失败，使用占位符函数抛出错误
     if model is None or tokenizer is None:
-        logger.warning("[HANDLER] 模型或分词器为 None，返回错误信息。")
-        yield "错误：模型未能成功加载，请检查后台日志。"
-        return
+        model_load_error_placeholder()
+
+    # 将 Gradio 的 history 格式转换为 Hugging Face 模板所需的格式
+    # history 的格式: [[user, assistant], [user, assistant], ...]
+    # messages 的格式: [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
+    messages = []
+    for turn in history:
+        user_message, bot_message = turn
+        messages.append({"role": "user", "content": user_message})
+        messages.append({"role": "assistant", "content": bot_message})
+    messages.append({"role": "user", "content": message})
+    
+    logger.info(f"[HANDLER] 转换后的 messages 列表:\n{json.dumps(messages, indent=2, ensure_ascii=False)}")
 
     try:
         logger.info("[HANDLER] 正在应用聊天模板...")
@@ -56,52 +83,78 @@ def predict(messages: list):
         ).to(model.device)
         logger.info(f"[HANDLER] 模板应用成功，输入 token 数量: {prompt_tokens.shape[-1]}")
 
+        # 使用 TextIteratorStreamer 实现流式输出
         streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-
+        
+        # 定义生成参数
         generation_kwargs = {
             "input_ids": prompt_tokens,
             "streamer": streamer,
             "max_new_tokens": 1024,
             "do_sample": True,
             "temperature": 0.7,
-            "top_p": 0.9,
+            "top_p": 0.9
         }
 
-        logger.info("[HANDLER] 准备启动生成线程...")
+        # 在单独的线程中运行模型生成，以避免阻塞UI
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
         logger.info("[HANDLER] 生成线程已启动，开始从 streamer 中读取数据...")
 
+        # 从 streamer 中逐个 token 地 yield，实现实时流式效果
         buffer = ""
         token_count = 0
         for new_text in streamer:
             token_count += 1
-            logger.info(f"[STREAM] 正在生成第 {token_count} 个 token: {repr(new_text)}")
+            if "�" in new_text:  # 过滤掉解码不完全的特殊字符
+                continue
+            logger.debug(f"[STREAM] 正在生成第 {token_count} 个 token: {repr(new_text)}")
             buffer += new_text
             yield buffer
         
         logger.info(f"[HANDLER] Streamer 读取完毕，共生成 {token_count} 个 token。")
+        thread.join() # 确保线程执行完毕
 
     except Exception as e:
         logger.error(f"[HANDLER] 在推理过程中发生错误: {e}", exc_info=True)
-        yield "抱歉，处理您的请求时遇到了一个内部错误。"
+        # 使用 gr.Error 在界面上优雅地显示错误信息
+        raise gr.Error(f"抱歉，处理您的请求时遇到了一个内部错误: {e}")
     
-    end_time = time.time()
-    logger.info(f"--- [END] predict 函数结束，总耗时: {end_time - start_time:.2f} 秒 ---")
+    finally:
+        end_time = time.time()
+        logger.info(f"--- [END] predict 函数结束，总耗时: {end_time - start_time:.2f} 秒 ---")
+
+# --- 3. 创建并配置Gradio界面 ---
 
-# --- 3. 创建Gradio界面 (已优化) ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# 你的自定义Qwen模型聊天机器人\n## 模型: {MODEL_ID}")
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css="footer {visibility: hidden}") as demo:
+    gr.Markdown(f"# 你的自定义Qwen模型聊天机器人\n## 模型: `{MODEL_ID}`")
     
+    # gr.ChatInterface 是 Gradio 5.x 中构建聊天机器人的推荐方式
     chat_interface = gr.ChatInterface(
         fn=predict,
-        title="聊天机器人",
-        description="向你的微调Qwen模型提问吧！",
-        examples=[["你好，你是谁？"], ["用Python写一个快速排序算法"]],
-        type="messages" # <-- 【最关键的优化】告诉 Gradio 使用新的 OpenAI 格式
+        # Gradio 5.x 的 `fn` 自动接收 message 和 history，无需手动管理状态
+        chatbot=gr.Chatbot(
+            height=600, 
+            show_copy_button=True,
+            avatar_images=(None, "https://s2.loli.net/2024/07/17/iPqD3uVgW9eBkbT.png") # (user, bot)
+        ),
+        title="Qwen 大模型聊天室",
+        description="向你的微调Qwen模型提问吧！这是一个流式输出的例子。",
+        examples=[
+            ["你好，你是谁？"],
+            ["用 Python 写一个快速排序算法。"],
+            ["解释一下什么是大型语言模型（LLM）。"]
+        ],
+        submit_btn="发送",
+        retry_btn="🔄 重试",
+        undo_btn="↩️ 撤销",
+        clear_btn="🗑️ 清除"
     )
 
 # --- 4. 启动应用 ---
-logger.info("准备启动 Gradio 应用...")
-demo.queue().launch()
-logger.info("Gradio 应用已启动。")
+
+if __name__ == "__main__":
+    logger.info("准备启动 Gradio 应用...")
+    # 使用 queue() 实现请求排队，concurrency_count 控制并发数
+    demo.queue(concurrency_count=2).launch(server_name="0.0.0.0", server_port=7860)
+    logger.info("Gradio 应用已启动。")