melbinjp/DocQA

1import datetime2import re3 4import faiss5import numpy as np6from rank_bm25 import BM25Okapi7 8# How many candidates each retriever contributes before fusion. Wider than the9# k finally returned, because the point of fusion is to let a chunk that one10# retriever ranked eighth and the other ranked second come out near the top.11CANDIDATE_POOL = 3012 13# The dense side searches vectors, and a long chunk owns several of them, so a14# pool of 30 vectors can collapse to six or seven distinct chunks. This is the15# vector pool, sized so the number of distinct chunks reaching fusion stays in16# the same range as the lexical side's.17DENSE_VECTOR_POOL = 12018 19# The constant in reciprocal rank fusion, from Cormack et al. 60 is the value20# the paper uses and the one every implementation inherits; it damps the21# difference between rank 1 and rank 2 so a single confident retriever cannot22# monopolise the result.23RRF_K = 6024 25_TOKEN = re.compile(r"[a-z0-9]+")26 27 28def _tokenize(text: str) -> list[str]:29    """Words and numbers, lowercased. Numbers matter here more than usual: the30    facts that dense retrieval misses in this corpus are table cells."""31    return _TOKEN.findall(text.lower())32 33class RAGSession:34    """35    Manages the RAG process for a single, isolated user session in memory.36 37    Each instance of this class handles the data for one ingested document,38    including its text chunks, embeddings, and a FAISS index for searching.39    It also tracks its last access time for automatic cleanup.40    """41    def __init__(self, source: str, embedding_model):42        self.source = source43        self.last_accessed = datetime.datetime.now()44        self.embedding_model = embedding_model45 46        d_model = self.embedding_model.get_sentence_embedding_dimension()47 48        # Inner product over L2-normalised vectors, which is cosine similarity.49        #50        # This was IndexFlatL2 with unnormalised vectors, and the score reported51        # to the user was `1 / (1 + l2_distance)`. That is not a similarity: it52        # is an unbounded distance squashed into (0, 1] with no meaning at53        # either end. Measured on the live Space, the five correctly retrieved54        # passages for a question about multi-head attention scored 0.081 down55        # to 0.076, and the UI rendered them as "Confidence: 8.1%" beside56        # answers that were right. Cosine puts the same passages near 1.57        #58        # It also ranks better. Untitled L2 distance is sensitive to vector59        # magnitude, and magnitude tracks chunk length more than relevance, so60        # long chunks were being penalised for being long.61        self.index = faiss.IndexFlatIP(d_model)62 63        # In-memory store for the actual text chunks. This is NO LONGER parallel64        # to the FAISS index: one chunk can own several vectors, so a vector id65        # is mapped through self.vector_parent to get here.66        self.chunks = []67 68        # vector id -> index into self.chunks.69        #70        # A long chunk is indexed once whole and once per window across it, so a71        # fact buried inside a chunk about something else is still reachable. The72        # RAG paper's "21M documents" sat in a 1435-character chunk that opens73        # "To estimate the probability of an hypothesis y"; three differently74        # worded questions about index size all failed to retrieve it, because75        # the chunk's vector is about marginalising over documents. Whichever76        # vector matches, the parent chunk is what is returned, read and cited.77        self.vector_parent = []78 79        # parent index -> one of its vector ids, so a chunk found only by the80        # lexical retriever can still have its cosine recovered for display.81        self.parent_vector = {}82 83        # Lexical half of the retriever, rebuilt whenever chunks are added.84        #85        # Dense embeddings are the wrong tool for part of a document and it is a86        # measurable part. Asked "how does the parameter count of the big model87        # compare to the base model", where the answer is a column headed88        # `params x10^6` holding 65 and 213, the dense retriever pulled the prose89        # of that page and never the table, and the API answered that it had no90        # such information twice: once with 500-character chunks and again with91        # 1500-character chunks and the table rendered as a grid. The chunk92        # existed and was correct both times. A block of bare numbers simply has93        # no embedding near "parameter count".94        #95        # `params` is right there as a literal token, which is what BM25 is for.96        self.bm25 = None97        self.tokenized = []98 99        # Parallel to self.chunks: the page each chunk came from, or None for100        # formats without pages. Kept as a separate list rather than a dict so a101        # FAISS vector id indexes both without a lookup that could disagree.102        self.pages = []103 104    def ingest(self, text_chunks: list[str], embeddings: np.ndarray,105               pages: list | None = None, vector_parents: list | None = None):106        """107        Processes and ingests text chunks and their pre-computed embeddings108        into the session's RAG store.109 110        Args:111            text_chunks: A list of strings, where each string is a chunk of the112                         source document.113            embeddings: A numpy array of the embeddings for the text chunks.114        """115        if not text_chunks:116            return117 118        # FAISS requires a flat numpy array of float32. Normalising in place is119        # what turns the inner-product index into a cosine index; skip it and120        # every score is meaningless and the ranking is magnitude-biased.121        embeddings_float32 = np.ascontiguousarray(np.array(embeddings, dtype='float32'))122        faiss.normalize_L2(embeddings_float32)123 124        # One vector per row of `embeddings`, each owned by a chunk. Without an125        # explicit mapping the relationship is one to one, which is what a caller126        # that knows nothing about windows should get.127        if vector_parents is None:128            vector_parents = list(range(len(text_chunks)))129        if len(vector_parents) != len(embeddings_float32):130            raise ValueError(131                f"{len(embeddings_float32)} vectors but {len(vector_parents)} owners; "132                "a citation would name the wrong chunk"133            )134 135        base_vector = self.index.ntotal136        base_chunk = len(self.chunks)137 138        # Add the new embeddings to the FAISS index.139        self.index.add(embeddings_float32)140 141        for offset, parent in enumerate(vector_parents):142            absolute_parent = base_chunk + parent143            self.vector_parent.append(absolute_parent)144            self.parent_vector.setdefault(absolute_parent, base_vector + offset)145 146        # Store the corresponding text chunks, and the page each came from. The147        # two lists must stay the same length or a citation would name the wrong148        # page, which is worse than naming none.149        self.chunks.extend(text_chunks)150        if pages is None:151            pages = [None] * len(text_chunks)152        if len(pages) != len(text_chunks):153            raise ValueError(154                f"{len(text_chunks)} chunks but {len(pages)} pages; a citation would be wrong"155            )156        self.pages.extend(pages)157 158        # Rebuilt rather than updated: BM25 scoring depends on corpus-wide159        # document frequencies and average length, so an index built over the160        # first document would score the second against the wrong statistics.161        self.tokenized.extend(_tokenize(c) for c in text_chunks)162        self.bm25 = BM25Okapi(self.tokenized) if self.tokenized else None163 164        print(f"Session ingested {self.index.ntotal} chunks.")165 166    async def query(self, query_text: str, k: int = 8) -> list[dict]:167        """168        Performs a similarity search against the session's document chunks.169 170        Args:171            query_text: The user's question.172            k: The number of top results to retrieve.173 174        Returns:175            A list of dictionaries, each containing the 'text' of a relevant176            chunk and its similarity 'score'.177        """178        import asyncio179        if self.index.ntotal == 0:180            return []181 182        # Embed the query.183        try:184            query_embedding_raw = await asyncio.wait_for(185                asyncio.to_thread(self.embedding_model.encode, [query_text], convert_to_numpy=True),186                timeout=30.0187            )188            query_embedding = np.ascontiguousarray(query_embedding_raw.astype('float32'))189            faiss.normalize_L2(query_embedding)190        except asyncio.TimeoutError:191            raise TimeoutError("Embedding generation for query timed out.")192 193        # Search the index. With a normalised inner-product index, `scores` are194        # cosine similarities in [-1, 1], already sorted high to low.195        pool = min(CANDIDATE_POOL, len(self.chunks))196        vector_pool = min(DENSE_VECTOR_POOL, self.index.ntotal)197        try:198            similarities, indices = await asyncio.wait_for(199                asyncio.to_thread(self.index.search, query_embedding, vector_pool),200                timeout=30.0201            )202        except asyncio.TimeoutError:203            raise TimeoutError("FAISS search for query timed out.")204 205        # Collapse vectors to the chunks that own them, keeping the best rank206        # and the best cosine each chunk achieved. Several windows of one chunk207        # can all match; that is one result, not four.208        dense_ranked, cosine = [], {}209        for position, vector_id in enumerate(indices[0]):210            if vector_id == -1:211                continue212            parent = self.vector_parent[int(vector_id)]213            similarity = float(similarities[0][position])214            if parent not in cosine or similarity > cosine[parent]:215                cosine[parent] = similarity216            if parent not in dense_ranked:217                dense_ranked.append(parent)218 219        # Lexical ranking, which is over whole chunks already.220        lexical_ranked = []221        if self.bm25 is not None:222            bm25_scores = await asyncio.to_thread(self.bm25.get_scores, _tokenize(query_text))223            lexical_ranked = [int(i) for i in np.argsort(bm25_scores)[::-1][:pool]224                              if bm25_scores[i] > 0]225 226        # Reciprocal rank fusion. Ranks rather than scores, because a cosine and227        # a BM25 score are not on the same scale and never will be; normalising228        # them against each other would be inventing a relationship.229        fused: dict[int, float] = {}230        for ranking in (dense_ranked, lexical_ranked):231            for rank, parent in enumerate(ranking):232                fused[parent] = fused.get(parent, 0.0) + 1.0 / (RRF_K + rank + 1)233 234        order = sorted(fused, key=lambda p: fused[p], reverse=True)[:k]235 236        results = []237        for parent in order:238            # The reported score stays a cosine, so it means one thing wherever239            # it appears. A chunk only the lexical side found has no cosine yet,240            # so it is recovered from one of that chunk's stored vectors rather241            # than left blank or filled with a fused rank score that would look242            # like a similarity and not be one.243            score = cosine.get(parent)244            if score is None:245                try:246                    vec = self.index.reconstruct(self.parent_vector[parent])247                    score = float(np.dot(query_embedding[0], vec))248                except Exception:249                    score = 0.0250 251            results.append({252                "text": self.chunks[parent],253                "score": float(score),254                "page": self.pages[parent] if parent < len(self.pages) else None,255            })256 257        return results258 259    def touch(self):260        """Updates the last_accessed timestamp to the current time."""261        self.last_accessed = datetime.datetime.now()262