| 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): |
| |
| 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() |
