Create rag.py

rag.py

@@ -0,0 +1,37 @@
+from sentence_transformers import SentenceTransformer
+import faiss
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+
+embed_model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
+
+
+index = faiss.read_index("index/family.faiss")
+
+
+with open("index/blocks.txt", "r", encoding="utf-8") as f:
+raw = f.read()
+blocks = raw.split("===")
+
+
+llm_name = "Qwen/Qwen2.5-1.5B"
+tokenizer = AutoTokenizer.from_pretrained(llm_name)
+model = AutoModelForCausalLM.from_pretrained(llm_name)
+
+
+SYSTEM_PROMPT = "你是冯氏家谱机器人，只能根据给定资料回答，没有记载就说未记载。"
+
+
+def ask(question):
+q_emb = embed_model.encode([question], normalize_embeddings=True)
+_, I = index.search(q_emb, 1)
+context = blocks[I[0][0]]
+
+
+prompt = f"{SYSTEM_PROMPT}\n\n【家谱资料】\n{context}\n\n【问题】{question}\n回答："
+
+
+inputs = tokenizer(prompt, return_tensors="pt")
+outputs = model.generate(**inputs, max_new_tokens=200)
+return tokenizer.decode(outputs[0], skip_special_tokens=True)