tensorsoft/Mini-RAG-Chat-With-Your-Files-CPU-Only

1from sentence_transformers import SentenceTransformer2import numpy as np3from utils import chunk_text, cosine_similarity4 5class TextRetriever:6    def __init__(self):7        self.model = SentenceTransformer('all-MiniLM-L6-v2')8        self.chunks = []9        self.embeddings = []10 11    def add_document(self, text: str):12        """Split text into chunks and store their embeddings."""13        self.chunks = chunk_text(text, chunk_size=200)14        self.embeddings = self.model.encode(self.chunks, convert_to_numpy=True)15 16    def retrieve(self, query: str, top_k: int = 3):17        """Retrieve top_k relevant chunks for the query."""18        if not self.chunks:19            return []20        query_embedding = self.model.encode([query], convert_to_numpy=True)[0]21        similarities = [cosine_similarity(query_embedding, emb) for emb in self.embeddings]22        top_indices = np.argsort(similarities)[-top_k:][::-1]23        return [self.chunks[i] for i in top_indices]24