Spaces:
Sleeping
Sleeping
File size: 2,118 Bytes
bf4478b | 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 60 | import os
import faiss
import json
from sentence_transformers import SentenceTransformer
KNOWLEDGE_BASE_DIR = "./knowledge_base"
FAISS_INDEX_PATH = "./faiss_index.bin"
METADATA_PATH = "./faiss_metadata.json"
EMBEDDING_MODEL = "BAAI/bge-small-en-v1.5"
print(f"--- Loading Embedding Model: {EMBEDDING_MODEL} ---")
embedder = SentenceTransformer(EMBEDDING_MODEL)
text_chunks = []
def chunk_document(text):
paragraphs = text.split("\n\n")
return [p.strip() for p in paragraphs if len(p.strip()) > 20]
def initialize_vector_store():
global text_chunks
embedding_dim = embedder.get_sentence_embedding_dimension()
if os.path.exists(FAISS_INDEX_PATH) and os.path.exists(METADATA_PATH):
print("--- Loading existing FAISS index ---")
index = faiss.read_index(FAISS_INDEX_PATH)
with open(METADATA_PATH, "r") as f:
text_chunks = json.load(f)
return index
print("--- Building FAISS index from scratch ---")
os.makedirs(KNOWLEDGE_BASE_DIR, exist_ok=True)
all_text = ""
for filename in os.listdir(KNOWLEDGE_BASE_DIR):
if filename.endswith(".txt") or filename.endswith(".md"):
with open(os.path.join(KNOWLEDGE_BASE_DIR, filename), "r", encoding="utf-8") as f:
all_text += f.read() + "\n\n"
if not all_text.strip():
all_text = "Flow matching is a generative modeling framework."
text_chunks = chunk_document(all_text)
print(f"--- Embedding {len(text_chunks)} chunks ---")
embeddings = embedder.encode(text_chunks, convert_to_numpy=True)
index = faiss.IndexFlatL2(embedding_dim)
index.add(embeddings)
faiss.write_index(index, FAISS_INDEX_PATH)
with open(METADATA_PATH, "w") as f:
json.dump(text_chunks, f)
return index
def retrieve_context(index, query, top_k=3):
query_vector = embedder.encode([query], convert_to_numpy=True)
distances, indices = index.search(query_vector, k=top_k)
retrieved_chunks = [text_chunks[i] for i in indices[0] if i != -1]
return "\n\n---\n\n".join(retrieved_chunks) |