Update app.py

app.py

@@ -1,102 +1,82 @@
+
 import gradio as gr
-import requests
+from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
 import os
-import json
 
 # --- 配置 ---
-# 从Hugging Face Space的Secrets中获取API Token
-# 请确保在你的Space设置中添加了名为 "HF_TOKEN" 的Secret
-HF_TOKEN = os.getenv("HF_TOKEN")
-API_URL = "https://api-inference.huggingface.co/models/badanwang/teacher_basic_qwen3-0.6b"
+# 我们不再需要API Token，因为模型在本地运行
+MODEL_ID = "badanwang/teacher_basic_qwen3-0.6b"
+
+
+print("开始加载模型和分词器...")
+try:
+    # 确保使用 trust_remote_code=True，因为Qwen模型需要加载自定义代码
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype="auto", # 使用适合CPU的类型，如torch.float32
+        device_map="auto",  # 自动将模型加载到可用设备（这里是CPU）
+        trust_remote_code=True
+    )
+    print("模型和分词器加载成功！")
+except Exception as e:
+    print(f"模型加载失败: {e}")
+    # 如果模型加载失败，应用将无法工作，这里可以抛出异常或退出
+    raise gr.Error(f"关键错误：无法加载模型 {MODEL_ID}。错误信息: {e}")
+
 
 # --- 核心对话函数 ---
 def predict(message, history):
     """
-    主函数，用于与Hugging Face Inference API进行流式对话。
-    :param message: 用户当前发送的消息 (str)
-    :param history: 对话历史 (list of lists)，格式为 [[user_msg, assistant_msg], ...]
-    :return: 一个生成器 (generator)，逐字(token)返回模型的响应
+    主函数，使用加载到本地的模型进行流式对话。
     """
-    if not HF_TOKEN:
-        raise gr.Error("Hugging Face API Token 未配置！请在Space的Secrets中添加 HF_TOKEN。")
-
-    headers = {
-        "Authorization": f"Bearer {HF_TOKEN}",
-        "Content-Type": "application/json"
-    }
-
-    # 1. 格式化对话历史以符合API要求
-    # API需要一个包含所有对话的列表，格式为 {"role": "user", "content": "..."} 或 {"role": "assistant", "content": "..."}
+    # 1. 格式化对话历史
+    # Qwen的模板要求一个特殊的列表格式
     messages = []
     for turn in history:
         user_msg, assistant_msg = turn
         messages.append({"role": "user", "content": user_msg})
         messages.append({"role": "assistant", "content": assistant_msg})
-    
-    # 添加当前用户消息
     messages.append({"role": "user", "content": message})
+    
+    # 使用分词器的 apply_chat_template 方法来正确格式化输入
+    model_inputs = tokenizer.apply_chat_template(
+        messages, 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to(model.device) # 确保输入张量和模型在同一设备上
 
-    # 2. 构建API请求体
-    # 我们启用流式响应 (stream=True)
-    payload = {
-        "inputs": messages,
-        "parameters": {
-            "max_new_tokens": 2048,  # 根据需要调整
-            "temperature": 0.7,
-            "top_p": 0.95,
-            "repetition_penalty": 1.1,
-            "return_full_text": False,
-        },
-        "stream": True
-    }
+    # 2. 设置流式输出
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
 
-    # 3. 发送流式请求并处理响应
+    # 3. 在一个单独的线程中运行生成，以避免阻塞UI
+    generation_kwargs = dict(
+        inputs=model_inputs,
+        streamer=streamer,
+        max_new_tokens=2048,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+    )
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    # 4. 从streamer中yield每个新生成的token
     full_response = ""
-    try:
-        # 使用 requests 发送POST请求，并设置 stream=True
-        with requests.post(API_URL, headers=headers, json=payload, stream=True, timeout=120) as response:
-            # 检查HTTP响应状态码
-            response.raise_for_status() 
-            
-            # 逐行读取流式响应
-            for line in response.iter_lines():
-                if line:
-                    # 流式响应通常以 "data:" 开头，后跟一个JSON对象
-                    decoded_line = line.decode('utf-8')
-                    if decoded_line.startswith("data:"):
-                        try:
-                            # 解析JSON
-                            json_data = json.loads(decoded_line[5:])
-                            # 提取token文本
-                            token = json_data.get("token", {}).get("text", "")
-                            if token:
-                                full_response += token
-                                yield full_response
-                        except json.JSONDecodeError:
-                            # 忽略无法解析的行
-                            continue
-                            
-    except requests.exceptions.RequestException as e:
-        print(f"API请求错误: {e}")
-        yield f"抱歉，与模型API通信时发生错误: {e}"
-    except Exception as e:
-        print(f"发生未知错误: {e}")
-        yield f"抱歉，发生了一个未知错误: {e}"
+    for new_text in streamer:
+        full_response += new_text
+        yield full_response
 
 # --- 创建并启动Gradio界面 ---
-
-# 使用gr.ChatInterface，它为聊天机器人提供了完整的UI
-# fn=predict 指定了处理逻辑的函数
-# streaming=True 告诉Gradio我们的函数是流式的（使用yield）
-# Gradio 4.44.1中，ChatInterface会自动处理stream参数，我们只需确保函数是生成器
 demo = gr.ChatInterface(
     fn=predict,
-    title="小Q老师 - 基础问答",
-    description="与 badanwang/teacher_basic_qwen3-0.6b 模型进行流式对话。直接输入问题开始。",
+    title="小Q老师 - 基础问答 (本地加载)",
+    description=f"直接在Space中运行 {MODEL_ID} 模型进行流式对话。CPU推理可能较慢，请耐心等待。",
     examples=[["你好"], ["请用python写一个快速排序算法"], ["给我讲个笑话吧"]],
     cache_examples=False,
 )
 
 if __name__ == "__main__":
-    # demo.launch(share=True) # 如果在本地运行并需要分享链接
-    demo.launch() # 在Hugging Face Spaces上运行时使用
+    demo.launch()