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("向量库构建完成")