RAGbot / build_index.py
fdbw's picture
Rename app.py to build_index.py
8f20187 verified
Raw
History Blame Contribute Delete
659 Bytes
from sentence_transformers import SentenceTransformer
import faiss
import os
model = SentenceTransformer("BAAI/bge-small-zh-v1.5")
with open("data/feng_family.txt", "r", encoding="utf-8") as f:
text = f.read()
blocks = [b.strip() for b in text.split("--- PERSON ---") if b.strip()]
embeddings = model.encode(blocks, normalize_embeddings=True)
index = faiss.IndexFlatIP(embeddings.shape[1])
index.add(embeddings)
os.makedirs("index", exist_ok=True)
faiss.write_index(index, "index/family.faiss")
with open("index/blocks.txt", "w", encoding="utf-8") as f:
for b in blocks:
f.write(b.replace("\n", " ") + "\n===\n")
print("向量库构建完成")