from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig
)
from peft import PeftModel
import torch
import gradio as gr
import os

# === 配置 ===
BASE_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
ADAPTER_PATH = "./"  # 因为模型文件就在 Space 根目录

print("🚀 正在加载基础模型...")

# 加载基础模型（自动使用 CPU）
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_NAME,
    torch_dtype=torch.float16,  # 即使 CPU 也建议用 float16 减少内存
    device_map="cpu",           # 明确指定 CPU
    trust_remote_code=False,
)

print("🔧 正在加载 LoRA 适配器...")

# 加载你微调的 LoRA 权重
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
model.eval()  # 切换到推理模式

print("🔤 正在加载分词器...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_PATH)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# === 推理函数 ===
def generate_response(prompt: str, history=None):
    try:
        # 构造输入（可根据你训练时的格式调整）
        input_text = prompt.strip()
        
        inputs = tokenizer(
            input_text,
            return_tensors="pt",
            truncation=True,
            max_length=256,
            padding=True
        ).to("cpu")
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=128,
                do_sample=True,
                temperature=0.7,
                top_p=0.9,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # 去掉输入部分，只返回生成内容
        if response.startswith(input_text):
            response = response[len(input_text):].strip()
        
        return response or "抱歉，我无法回答这个问题。"
    
    except Exception as e:
        return f"❌ 推理出错: {str(e)}"

# === Gradio 界面 ===
with gr.Blocks(title="冯氏家谱助手") as demo:
    gr.Markdown("# 🧬 冯氏家族知识问答\n基于 TinyLlama 微调的家谱 AI 助手")
    chatbot = gr.ChatInterface(
        fn=generate_response,
        examples=["冯国璋的字辈是什么？", "冯玉祥生于哪一年？", "冯家第几代是‘国’字辈？"],
        title="冯氏家谱助手"
    )

# 启动
demo.launch()