import gradio as gr
from transformers import pipeline

print("正在初始化纯 Python 推理引擎（免编译方案）...")

try:
    # 使用 Hugging Face 官方的 pipeline 直接加载 EdgeAI 模型的 GGUF 文件
    pipe = pipeline(
        "text-generation", 
        model="eoinedge/edgeai-docs-qwen2.5-coder-0.5b-lora", 
        model_kwargs={"gguf_file": "edgeai-docs-qwen2.5-coder-0.5b-lora.Q4_K_M.gguf"}
    )
    print("EdgeAI 模型加载成功！")
except Exception as e:
    pipe = None
    print(f"模型加载失败: {e}")

def predict(message, history):
    if pipe is None:
        return "服务器模型加载失败，请检查模型路径。"
    
    # 🧠 自动构建带有聊天记忆的上下文 Prompt
    prompt = ""
    for user_msg, ai_msg in history:
        prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n<|im_start|>assistant\n{ai_msg}<|im_end|>\n"
    
    # 加上当前的最新提问
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    
    # 开始推理
    outputs = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)
    
    # 提取新生成的文本
    generated_text = outputs[0]["generated_text"]
    response = generated_text[len(prompt):].replace("<|im_end|>", "").strip()
    return response

# 构筑高颜值聊天网页
gr.ChatInterface(
    fn=predict, 
    title="🤖 EdgeAI 网页编程助手 (纯净免编译版)",
    description="已完美集成 Edge Impulse 官方文档微调模型，支持上下文对话记忆，由 Hugging Face 独立沙盒托管。"
).launch()