import gradio as gr
from llama_cpp import Llama
import os

# 从 Hugging Face Model Hub 下载 GGUF 模型文件
# 这里以 TheBloke 量化的大佬提供的模型为例
model_repo = "Qwen/Qwen3-4B-GGUF"
model_file = "Qwen3-4B-Q8_0.gguf" 
# 也可以选择更小参数的量化版本，如 q8_0, q5_0, q4_0 等，数字越小，模型越小，精度越低

# 初始化 Llama 模型
# 首次运行会自动下载模型，需要等待一段时间
llm = Llama(
    model_path=f"https://huggingface.co/{model_repo}/resolve/main/{model_file}",
    n_ctx=2048,  # 上下文长度
    n_threads=2, # 线程数
    verbose=False
)

def predict(message, history):
    # 1. 构建聊天历史格式
    system_message = {"role": "system", "content": "You are a helpful assistant."}
    messages = [system_message]
    
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})
    
    # 2. 生成回复
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=512,
        temperature=0.7,
        stop=["<|im_end|>"] # Qwen 模型的停止词
    )
    
    # 3. 提取回复内容
    assistant_reply = response['choices'][0]['message']['content']
    return assistant_reply

# 创建界面
gr.ChatInterface(
    fn=predict,
    title="Qwen3-4B-GGUF (GGUF量化版)",
    description="使用 llama.cpp 在 CPU 上高效运行 Qwen 模型。首次加载需下载模型，请耐心等待。",
).launch(server_name="0.0.0.0", server_port=7860)