import gradio as gr
import pickle
import faiss
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import numpy as np
import os

# --- 配置 ---
LLM_MODEL = "Qwen/Qwen1.5-0.5B-Chat" 
EMBEDDING_MODEL = "BAAI/bge-small-zh-v1.5"
FAISS_PATH = 'data/family_rag.faiss'
CHUNKS_PATH = 'data/family_chunks.pkl'

device = "cpu"

def rag_answer(query):
    # 【检测点 1】检查文件是否存在
    if not os.path.exists(FAISS_PATH) or not os.path.exists(CHUNKS_PATH):
        return "❌ 错误：在 data 文件夹下未找到索引文件。请检查是否已上传 family_rag.faiss 和 family_chunks.pkl。", "无资料"

    try:
        # 加载资源
        embedder = SentenceTransformer(EMBEDDING_MODEL, device=device)
        index = faiss.read_index(FAISS_PATH)
        with open(CHUNKS_PATH, 'rb') as f:
            corpus_chunks = pickle.load(f)
        
        tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
        model = AutoModelForCausalLM.from_pretrained(LLM_MODEL, torch_dtype=torch.float32)

        # 搜索
        query_embedding = embedder.encode(query, convert_to_tensor=False).astype('float32')
        _, I = index.search(np.expand_dims(query_embedding, axis=0), 1)
        
        context = corpus_chunks[I[0][0]] if I[0][0] != -1 else "未找到资料"
        
        # 生成回答
        prompt = f"资料：{context}\n问题：{query}\n答案："
        inputs = tokenizer([prompt], return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=100, do_sample=False)
        
        res = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
        return res.split("答案：")[-1].strip(), context

    except Exception as e:
        return f"❌ 运行中出错：{str(e)}", "无资料"

# --- 极简界面 ---
demo = gr.Interface(
    fn=rag_answer,
    inputs=gr.Textbox(label="输入人名"),
    outputs=[gr.Textbox(label="回答"), gr.Textbox(label="参考资料")],
    title="冯氏家谱查询系统"
)

demo.launch()