Spaces:
Sleeping
Sleeping
| 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) |