code

app.py

@@ -1,20 +1,15 @@
 import spaces
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from llama_cpp import Llama
 import os
 
-if not os.path.exists("./phi-4"):
-    os.system('huggingface-cli download matteogeniaccio/phi-4 --local-dir ./phi-4 --include "phi-4/*"')
-
-import transformers
-
-# 初始化pipeline
-pipeline = transformers.pipeline(
-    "text-generation",
-    model="./phi-4/phi-4",
-    model_kwargs={"torch_dtype": "auto"},
-    device_map="auto",
+# 初始化LLM
+llm = Llama.from_pretrained(
+    repo_id="matteogeniaccio/phi-4",
+    filename="phi-4-Q4_K_M.gguf",
+    verbose=True,
+    main_gpu=0,
+    n_gpu_layers=-1
 )
 
 # 响应函数
@@ -36,19 +31,19 @@ def respond(
             messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
 
-    # 生成响应
-    outputs = pipeline(
-        messages,
-        max_new_tokens=max_tokens,
+    # 使用llama-cpp-python的方式生成响应
+    response = llm.create_chat_completion(
+        messages=messages,
+        max_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        do_sample=(temperature > 0),
+        stream=True  # 启用流式输出
     )
-    response = outputs[0]["generated_text"]
 
     # 返回流式响应
-    for token in response:
-        yield token
+    for chunk in response:
+        if chunk and chunk.get("choices") and chunk["choices"][0].get("delta", {}).get("content"):
+            yield chunk["choices"][0]["delta"]["content"]
 
 # Gradio 界面
 demo = gr.ChatInterface(