Update app.py

app.py

@@ -1,131 +1,94 @@
-import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
-import logging
-import time
-import json
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import os
 
+# --- 1. 配置与模型加载 ---
 
-# 日志记录配置
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# 从环境变量或默认值加载模型ID，增加灵活性
+# 从环境变量或默认值加载模型ID
 MODEL_ID = os.getenv("MODEL_ID", "badanwang/teacher_basic_qwen3-0.6b")
-
-logger.info("===== Application Startup =====")
-logger.info(f"正在加载模型和分词器: {MODEL_ID}")
-
-try:
-    # 推荐使用 trust_remote_code=True 以确保所有模型组件正确加载
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        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. 核心推理函数 (无变动) ---
-
-def predict(message: str, history: list[list[str]]):
-    start_time = time.time()
-    logger.info("\n--- [START] 进入 predict 函数 ---")
-    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:
-        model_load_error_placeholder()
-
+print(f"正在加载模型: {MODEL_ID}")
+
+# 加载分词器和模型
+# trust_remote_code=True 是加载Qwen等模型所必需的
+# device_map="auto" 会自动将模型分配到可用的硬件上（如GPU）
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype="auto",
+    device_map="auto",
+    trust_remote_code=True
+)
+print("模型加载成功！")
+
+
+# --- 2. 核心推理函数 (API) ---
+
+def get_response(prompt: str, history: list[list[str]] = None):
+    """
+    一个简单的函数，用于与模型进行单次对话。
+
+    Args:
+        prompt (str): 用户当前输入的问题。
+        history (list[list[str]], optional): 对话历史，格式为 [[user_msg_1, bot_msg_1], ...]。默认为 None。
+
+    Returns:
+        str: 模型生成的回复。
+    """
+    if history is None:
+        history = []
+
+    # 1. 构建消息列表
     messages = []
-    for turn in history:
-        user_message, bot_message = turn
+    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": message})
-
-    logger.info(f"[HANDLER] 转换后的 messages 列表:\n{json.dumps(messages, indent=2, ensure_ascii=False)}")
-
-    try:
-        logger.info("[HANDLER] 正在应用聊天模板...")
-        prompt_tokens = tokenizer.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=True,
-            return_tensors="pt"
-        ).to(model.device)
-        logger.info(f"[HANDLER] 模板应用成功，输入 token 数量: {prompt_tokens.shape[-1]}")
-
-        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
-        }
-
-        thread = Thread(target=model.generate, kwargs=generation_kwargs)
-        thread.start()
-        logger.info("[HANDLER] 生成线程已启动，开始从 streamer 中读取数据...")
-
-        buffer = ""
-        token_count = 0
-        for new_text in streamer:
-            token_count += 1
-            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)
-        raise gr.Error(f"抱歉，处理您的请求时遇到了一个内部错误: {e}")
-
-    finally:
-        end_time = time.time()
-        logger.info(f"--- [END] predict 函数结束，总耗时: {end_time - start_time:.2f} 秒 ---")
-
-# --- 3. 创建并配置Gradio界面 (已优化) ---
+    messages.append({"role": "user", "content": prompt})
+
+    # 2. 应用聊天模板并进行分词
+    # 这是与聊天模型正确交互的关键步骤
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    # 3. 生成回复
+    # 这是一个阻塞式调用，会等待模型生成完毕
+    outputs = model.generate(
+        input_ids,
+        max_new_tokens=1024,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
+    
+    # 4. 解码生成的文本
+    # `outputs[0]` 包含了输入的token和新生成的token，我们需要切片只获取新生成的部分
+    response_ids = outputs[0][input_ids.shape[-1]:]
+    response_text = tokenizer.decode(response_ids, skip_special_tokens=True)
 
-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}`")
+    return response_text
 
-    chat_interface = gr.ChatInterface(
-        fn=predict,
-        chatbot=gr.Chatbot(
-            height=600,
-            show_copy_button=True,
-            avatar_images=(None, "https://s2.loli.net/2024/07/17/iPqD3uVgW9eBkbT.png")
-        ),
-        title="Qwen 大模型聊天室",
-        description="向你的微调Qwen模型提问吧！这是一个流式输出的例子。",
-        examples=[
-            ["你好，你是谁？"],
-            ["用 Python 写一个快速排序算法。"],
-            ["解释一下什么是大型语言模型（LLM）。"]
-        ],
-        submit_btn="发送",
-    )
+# --- 3. 使用示例 ---
 
 if __name__ == "__main__":
-    logger.info("准备启动 Gradio 应用...")
-    # .queue() 对于处理多个并发用户至关重要
-    # 在Hugging Face Spaces上部署时，share=True 不是必需的，但有助于本地测试
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)
-    logger.info("Gradio 应用已启动。")
+    # 示例1: 单轮对话
+    print("\n--- 示例 1: 单轮对话 ---")
+    question1 = "你好，你是谁？"
+    print(f"用户: {question1}")
+    answer1 = get_response(question1)
+    print(f"模型: {answer1}")
+
+    # 示例2: 多轮对话
+    print("\n--- 示例 2: 多轮对话 ---")
+    # 首先，定义一个对话历史
+    chat_history = [
+        ["用Python写一个快速排序", "当然，这是快速排序的Python实现：\n```python\ndef quick_sort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quick_sort(left) + middle + quick_sort(right)\n\nprint(quick_sort())\n```"]
+    ]
+    question2 = "很好，你能解释一下它的工作原理吗？"
+    print(f"历史: {chat_history}")
+    print(f"用户: {question2}")
+    # 调用时传入历史记录
+    answer2 = get_response(question2, history=chat_history)
+    print(f"模型: {answer2}")