File size: 2,107 Bytes
669172d 56c1e74 669172d 5fe8128 5b9232f 669172d 5b9232f 5fe8128 669172d 56c1e74 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | 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() |