Spaces:
Sleeping
Sleeping
| """ | |
| rag.py — Core RAG engine: chunk → embed → index → search. | |
| Works with Hebrew, Russian, and English out of the box. | |
| """ | |
| import json | |
| import os | |
| import pickle | |
| import numpy as np | |
| import faiss | |
| from sentence_transformers import SentenceTransformer | |
| # ── Config ────────────────────────────────────────────────── | |
| MODEL_NAME = "intfloat/multilingual-e5-base" # great Hebrew support | |
| INDEX_DIR = "index" | |
| CHUNK_SIZE = 400 # characters per chunk | |
| CHUNK_OVERLAP = 80 # overlap between chunks | |
| class TextChunker: | |
| """Split long texts into overlapping chunks with metadata.""" | |
| def __init__(self, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP): | |
| self.chunk_size = chunk_size | |
| self.overlap = overlap | |
| def chunk_text(self, text: str, metadata: dict = None) -> list[dict]: | |
| """Split text into overlapping chunks, preserving sentence boundaries.""" | |
| if metadata is None: | |
| metadata = {} | |
| # Split by paragraphs first, then by sentences | |
| paragraphs = [p.strip() for p in text.split("\n") if p.strip()] | |
| chunks = [] | |
| current_chunk = "" | |
| for para in paragraphs: | |
| if len(current_chunk) + len(para) + 1 <= self.chunk_size: | |
| current_chunk += ("\n" + para if current_chunk else para) | |
| else: | |
| if current_chunk: | |
| chunks.append({ | |
| "text": current_chunk, | |
| **metadata, | |
| }) | |
| # Start new chunk with overlap from previous | |
| if self.overlap > 0 and current_chunk: | |
| overlap_text = current_chunk[-self.overlap:] | |
| current_chunk = overlap_text + "\n" + para | |
| else: | |
| current_chunk = para | |
| if current_chunk: | |
| chunks.append({"text": current_chunk, **metadata}) | |
| return chunks | |
| def chunk_documents(self, documents: list[dict]) -> list[dict]: | |
| """Chunk a list of extracted documents.""" | |
| all_chunks = [] | |
| for doc in documents: | |
| text = doc["text"] | |
| metadata = {k: v for k, v in doc.items() if k != "text"} | |
| chunks = self.chunk_text(text, metadata) | |
| all_chunks.extend(chunks) | |
| return all_chunks | |
| class RAGIndex: | |
| """FAISS-based vector index with semantic search.""" | |
| def __init__(self, model_name: str = MODEL_NAME): | |
| print(f"🔄 Loading embedding model: {model_name}") | |
| self.model = SentenceTransformer(model_name) | |
| self.dimension = self.model.get_sentence_embedding_dimension() | |
| self.index = None | |
| self.chunks = [] | |
| def _embed(self, texts: list[str]) -> np.ndarray: | |
| """Embed texts with E5 prefix for better retrieval.""" | |
| # multilingual-e5 uses "query: " and "passage: " prefixes | |
| return self.model.encode(texts, normalize_embeddings=True, show_progress_bar=True) | |
| def build(self, chunks: list[dict]): | |
| """Build FAISS index from text chunks.""" | |
| self.chunks = chunks | |
| texts = [f"passage: {c['text']}" for c in chunks] | |
| print(f"🧮 Embedding {len(texts)} chunks...") | |
| embeddings = self._embed(texts) | |
| # Use IndexFlatIP for cosine similarity (embeddings are normalized) | |
| self.index = faiss.IndexFlatIP(self.dimension) | |
| self.index.add(embeddings.astype("float32")) | |
| print(f"✅ Index built: {self.index.ntotal} vectors, dim={self.dimension}") | |
| def search(self, query: str, top_k: int = 5) -> list[dict]: | |
| """Semantic search: returns top_k most relevant chunks.""" | |
| query_vec = self._embed([f"query: {query}"]) | |
| scores, indices = self.index.search(query_vec.astype("float32"), top_k) | |
| results = [] | |
| for score, idx in zip(scores[0], indices[0]): | |
| if idx == -1: | |
| continue | |
| chunk = self.chunks[idx].copy() | |
| chunk["score"] = float(score) | |
| results.append(chunk) | |
| return results | |
| def save(self, index_dir: str = INDEX_DIR): | |
| """Save index and chunks to disk.""" | |
| os.makedirs(index_dir, exist_ok=True) | |
| faiss.write_index(self.index, os.path.join(index_dir, "faiss.index")) | |
| with open(os.path.join(index_dir, "chunks.pkl"), "wb") as f: | |
| pickle.dump(self.chunks, f) | |
| print(f"💾 Index saved to {index_dir}/") | |
| def load(self, index_dir: str = INDEX_DIR): | |
| """Load index and chunks from disk.""" | |
| self.index = faiss.read_index(os.path.join(index_dir, "faiss.index")) | |
| with open(os.path.join(index_dir, "chunks.pkl"), "rb") as f: | |
| self.chunks = pickle.load(f) | |
| print(f"📂 Index loaded: {self.index.ntotal} vectors") | |
| # ── CLI ───────────────────────────────────────────────────── | |
| if __name__ == "__main__": | |
| import sys | |
| if len(sys.argv) > 1 and sys.argv[1] == "build": | |
| # Build index from extracted texts | |
| with open("data/extracted_texts.json", "r", encoding="utf-8") as f: | |
| documents = json.load(f) | |
| chunker = TextChunker() | |
| chunks = chunker.chunk_documents(documents) | |
| print(f"📝 {len(chunks)} chunks created") | |
| rag = RAGIndex() | |
| rag.build(chunks) | |
| rag.save() | |
| elif len(sys.argv) > 1 and sys.argv[1] == "search": | |
| query = " ".join(sys.argv[2:]) | |
| rag = RAGIndex() | |
| rag.load() | |
| results = rag.search(query) | |
| print(f"\n🔍 Query: {query}\n{'─' * 60}") | |
| for i, r in enumerate(results, 1): | |
| src = r.get("source", "?") | |
| score = r["score"] | |
| text = r["text"][:200] | |
| print(f"\n [{i}] (score: {score:.3f}) [{src}]") | |
| print(f" {text}...") | |
| else: | |
| print("Usage:") | |
| print(" python rag.py build # Build index from data/extracted_texts.json") | |
| print(" python rag.py search <query> # Search the index") | |