Update app.py

app.py

@@ -1,18 +1,22 @@
+import torch
 import re
 import gradio as gr
-from llama_cpp import Llama
-
-# 1. 加载 GGUF 量化模型（带 AVX2 优化后推理更快）
-llm = Llama.from_pretrained(
-    repo_id="Qwen/Qwen2.5-1.5B-Instruct-GGUF",
-    filename="*q4_k_m.gguf",
-    n_ctx=1024,       # 从 2048 降到 1024，大幅减少 KV cache 内存
-    n_threads=2,
-    n_batch=256,      # 从 512 降到 256
-    verbose=False,
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+
+# 1. 设置 CPU 线程数
+torch.set_num_threads(2)
+
+# 2. 加载模型（float32 兼容 CPU，0.5B 体积小速度快）
+model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.float32,
+    device_map="cpu",
 )
 
-# 2. 流式生成函数
+# 3. 流式生成函数
 def stream_response(message, history):
     messages = [{"role": "system", "content": "你是一个乐于助人的 AI 助手。"}]
     for h in history:
@@ -25,21 +29,35 @@ def stream_response(message, history):
                 messages.append({"role": "assistant", "content": h[1]})
     messages.append({"role": "user", "content": message})
 
-    # 3. 流式生成
+    model_inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+    ).to("cpu")
+
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+
+    generate_kwargs = {
+        **model_inputs,
+        "streamer": streamer,
+        "max_new_tokens": 512,
+        "do_sample": True,
+        "temperature": 0.7,
+        "pad_token_id": tokenizer.eos_token_id,
+    }
+
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
+
     partial_text = ""
-    for chunk in llm.create_chat_completion(
-        messages=messages,
-        max_tokens=512,
-        temperature=0.7,
-        stream=True,
-    ):
-        delta = chunk["choices"][0]["delta"]
-        if "content" in delta:
-            partial_text += delta["content"]
-
-            # --- 实时符号过滤（保留你原有的逻辑） ---
-            clean_display_text = re.sub(r'[\n\-\*\"""\"#]', '', partial_text)
-            yield clean_display_text
+    for new_text in streamer:
+        partial_text += new_text
+        clean_display_text = re.sub(r'[
+\-\*\"""""#]', '', partial_text)
+        yield clean_display_text
 
 # 4. 界面设置
 demo = gr.ChatInterface(
@@ -50,4 +68,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()