import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# 设置模型名称
model_id = "Qwen/Qwen3-1.7B-GGUF"

# 从 Hugging Face 加载模型和分词器
# 使用 `device_map="cpu"` 强制在 CPU 上运行，`torch_dtype=torch.float16` 可以节省内存
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto", # 如果是 CPU 环境，会自动 fallback 到 CPU
    trust_remote_code=True # Qwen 模型需要此选项
)

# 创建文本生成管道
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=512, # 生成的最大token数
)

# 定义聊天函数
def predict(message, history):
    """
    history 格式: [(用户消息1, 助手回复1), (用户消息2, 助手回复2), ...]
    """
    # 1. 将对话历史格式化为 Qwen 所需的聊天格式
    conversation = []
    for human, assistant in history:
        conversation.append({"role": "user", "content": human})
        conversation.append({"role": "assistant", "content": assistant})
    # 2. 加入当前用户的新消息
    conversation.append({"role": "user", "content": message})
    
    # 3. 应用聊天模板
    prompt = tokenizer.apply_chat_template(
        conversation,
        tokenize=False,
        add_generation_prompt=True # 添加让模型开始回复的提示
    )
    
    # 4. 生成回复
    outputs = pipe(prompt)
    response = outputs[0]['generated_text'][len(prompt):].strip() # 只提取新生成的部分
    
    return response

# 创建 Gradio 聊天界面
gr.ChatInterface(
    fn=predict,
    title="Qwen3-1.7B-GGUF 模型的LLM",
    description="这是一个基于 Qwen/Qwen3-1.7B-GGUF 模型的LLM。",
).launch(server_name="0.0.0.0", server_port=7860) # 在 Space 环境中必须这样设置