Update app.py

app.py

@@ -1,90 +1,165 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import os
 
-# --- 1. 配置与模型加载 ---
-# 假设运行环境的硬件资源是充足的。
+# --- 配置 ---
 MODEL_ID = os.getenv("MODEL_ID", "badanwang/teacher_basic_qwen3-0.6b")
-print(f"INFO: 正在加载模型: {MODEL_ID}")
+print(f"INFO: Application startup. Loading model: {MODEL_ID}")
 
-# 使用 try-except 来捕获任何可能的加载错误 (例如网络问题、模型名称错误等)
+# --- 1. 模型加载 (内置健壮的错误处理) ---
 try:
-    # 加载分词器和模型
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-    # device_map="auto" 会自动利用可用的硬件 (如 CPU 或 GPU)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        torch_dtype="auto",  # 自动选择最佳数据类型
+        torch_dtype="auto",
         device_map="auto",
         trust_remote_code=True
     )
-    print("INFO: 模型和分词器加载成功！")
+    print("INFO: Model and tokenizer loaded successfully!")
+    model_loaded = True
+except Exception as e:
+    print(f"FATAL: Failed to load model or tokenizer: {e}")
+    model_loaded = False
+    model_load_error = e
+
+# --- 2. 核心流式推理函数 ---
+def stream_predict(prompt: str, history: list[list[str]]):
+    """
+    一个生成器函数，用于流式生成对话。
+    它会逐步 (yield) 返回完整的对话历史。
+    """
+    if not model_loaded:
+        # 如果模型加载失败，则立即抛出错误
+        raise gr.Error(f"Model is not loaded. Please check logs. Error: {model_load_error}")
+
+    print(f"INFO: Received prompt: '{prompt}'")
     
-    # 将核心推理逻辑定义为一个函数
-    # 只有在模型成功加载后，这个函数才会被有效定义
-    def predict(prompt: str, history: list[list[str]]):
-        """
-        接收用户输入和对话历史，返回更新后的完整对话历史。
-        Gradio 会自动为这个函数创建 API 端点。
-        """
-        print(f"INFO: 收到API/UI请求: prompt='{prompt}'")
-        
-        # 1. 构建符合模型要求的消息列表
-        messages = []
-        for user_message, bot_message in history:
-            messages.append({"role": "user", "content": user_message})
-            messages.append({"role": "assistant", "content": bot_message})
-        messages.append({"role": "user", "content": prompt})
-
-        # 2. 应用聊天模板并进行分词
+    # 将历史记录和新提示转换为模型需要的格式
+    messages = []
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": prompt})
+
+    # 应用聊天模板
+    try:
         input_ids = tokenizer.apply_chat_template(
             messages,
             add_generation_prompt=True,
             tokenize=True,
             return_tensors="pt"
         ).to(model.device)
+    except Exception as e:
+        raise gr.Error(f"Error applying chat template: {e}")
 
-        # 3. 生成回复
-        # 使用简单的 .generate()，不加额外的采样参数以保持简洁
-        outputs = model.generate(input_ids, max_new_tokens=1024)
-        
-        # 4. 解码生成的文本，跳过输入的token
-        response_text = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True)
+    # 初始化 streamer 和生成线程
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+    
+    # 在独立线程中运行生成，防止阻塞UI
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 流式输出
+    try:
+        # 初始化一个空的字符串来存放助手的回复
+        assistant_response = ""
+        # 每次从streamer中获取一个新的文本片段
+        for new_text in streamer:
+            if not new_text:
+                continue
+            assistant_response += new_text
+            # 将当前用户输入和不断增长的助手回复组合成新的对话历史
+            # 然后使用 yield 返回，Gradio会用它来更新UI
+            yield history + [[prompt, assistant_response]]
         
-        print(f"INFO: 生成回复: {response_text}")
+        print("INFO: Streaming finished.")
+
+    except Exception as e:
+        print(f"ERROR: An error occurred during streaming: {e}")
+        raise gr.Error(f"An error occurred during generation: {e}")
+    finally:
+        # 确保线程结束
+        thread.join()
+
+
+# --- 3. Gradio Blocks 界面布局 ---
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), css="footer {visibility: hidden}") as demo:
+    gr.Markdown(f"# 流式对话机器人\n### 模型: `{MODEL_ID}`")
+    
+    # 使用 gr.State 来存储对话历史
+    # 这是实现多轮对话的关键
+    chatbot_state = gr.State([])
+    
+    # Chatbot 组件用于显示对话
+    chatbot_ui = gr.Chatbot(label="对话窗口", height=600)
+
+    with gr.Row():
+        # Textbox 用于用户输入
+        prompt_input = gr.Textbox(
+            show_label=False,
+            placeholder="请在这里输入您的问题...",
+            scale=4,
+        )
+        # Button 用于提交
+        submit_button = gr.Button("发送", variant="primary", scale=1)
+
+    # 清除按钮
+    clear_button = gr.Button("清除对话历史")
+
+    # --- 4. 事件处理逻辑 ---
+    
+    # 提交逻辑:
+    # 1. 点击"发送"按钮或在输入框按回车时触发
+    # 2. 调用 stream_predict 函数
+    # 3. 输入是用户输入框(prompt_input)和对话历史状态(chatbot_state)
+    # 4. 输出会实时更新聊天机器人界面(chatbot_ui)
+    # 5. 在函数开始前，将用户输入添加到聊天记录的末尾，并清空输入框
+    def on_submit(prompt, history):
+        # 将用户输入加入历史，形成 "用户: XXX" 的临时记录
+        return "", history + [[prompt, None]]
         
-        # 5. 更新并返回对话历史
-        history.append([prompt, response_text])
-        return history
+    prompt_input.submit(
+        on_submit, 
+        [prompt_input, chatbot_state], 
+        [prompt_input, chatbot_ui]
+    ).then(
+        stream_predict, 
+        [prompt_input, chatbot_state], 
+        chatbot_ui
+    )
 
-except Exception as e:
-    print(f"FATAL: 加载模型或分词器时发生致命错误: {e}")
-    # 如果模型加载失败，则定义一个专门用于报错的函数
-    # 这能确保Gradio界面依然可以启动，并向用户显示一个清晰的错误信息
-    def predict(*args, **kwargs):
-        raise gr.Error(f"模型未能加载，应用无法工作。请检查后台日志获取详细错误信息。错误: {e}")
-
-# --- 2. 创建并启动 Gradio 应用 ---
-# 使用 gr.Blocks 来自定义界面布局
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
-    gr.Markdown(f"## 模型聊天机器人\n当前模型: `{MODEL_ID}`")
+    submit_button.click(
+        on_submit, 
+        [prompt_input, chatbot_state], 
+        [prompt_input, chatbot_ui]
+    ).then(
+        stream_predict, 
+        [prompt_input, chatbot_state], 
+        chatbot_ui
+    )
     
-    # 定义聊天机器人组件和输入框
-    chatbot = gr.Chatbot(label="对话历史", height=600)
-    msg_input = gr.Textbox(label="在这里输入你的问题...", placeholder="例如：你好，你是谁？")
-    clear_button = gr.Button("清除对话")
-
-    # 设定组件的交互逻辑
-    # 当用户在输入框中按回车时，调用 predict 函数
-    msg_input.submit(predict, [msg_input, chatbot], chatbot)
-    # 当用户点击“清除对话”按钮时，清空聊天机器人组件
-    clear_button.click(lambda: [], None, chatbot)
-
-# --- 3. 启动应用并开放API ---
-print("INFO: 准备启动Gradio应用...")
-
-# .queue() 使应用能够处理多个排队的请求，并且在 4.29.0 版本中会自动开放API。
-# share=True 是解决CORS问题的关键。它会生成一个公开的、已配置好CORS的 .gradio.live 网址。
-# *** 已移除 'api_open=True' 参数以适配 gradio==4.29.0 ***
-demo.queue().launch(share=True)
+    # 清除逻辑:
+    # 点击按钮时，清空状态和UI
+    def on_clear():
+        return []
+
+    clear_button.click(on_clear, [], chatbot_state)
+    clear_button.click(on_clear, [], chatbot_ui)
+
+
+# --- 5. 启动应用 ---
+print("INFO: Preparing to launch Gradio app...")
+# .queue() 启用请求队列，对于流式应用是必需的
+# 在Hugging Face Spaces上, 无需 share=True, Gradio会自动处理
+demo.queue().launch()