tqdm

diff --git a/src/rag/__init__.py b/src/rag/__init__.py
index 79e5b4ca5e9e8f7919112fc8a3419405eff9f0ca..fe805a684af826221959284b7952c2e71022f6cb 100644 (file)
--- a/src/rag/__init__.py
+++ b/src/rag/__init__.py
@@ -1,8 +1,10 @@
 import logging
 import os
+from sys import stderr
 
 import fire
 import pydantic
+from tqdm import tqdm
 from pydantic_core import ValidationError
 
 from rag.chunking import FileType, chunk_file
@@ -39,7 +41,9 @@ class RAG:
         max_chunk_size: int = 2000,
     ) -> None:
         chunks = {}
-        for dirpath, _, files in os.walk(data):
+        for dirpath, _, files in tqdm(
+            list(os.walk(data)), desc="Collecting files"
+        ):
             for file in files:
                 path = dirpath + "/" + file
                 if path.endswith(".py"):
@@ -49,10 +53,13 @@ class RAG:
                 else:
                     continue
                 chunks.update(chunk_file(path, max_chunk_size, filetype))
-        bm25 = BM25.from_corpus(chunks, fetch_stopwords())
+        bm25 = BM25.from_corpus(
+            tqdm(chunks.items(), desc="Indexing files"), fetch_stopwords()
+        )
         try:
             os.makedirs(index, exist_ok=True)
             with open(index + "/index.json", "wb") as f:
+                print(f"Writing index to {index}...", file=stderr)
                 f.write(
                     storage_adapter.dump_json(
                         {
@@ -75,7 +82,10 @@ class RAG:
             with open(index + "/index.json") as f:
                 bm25: BM25[MinimalSource] = BM25.from_storage(
                     (key_adapter.validate_json(k), v)
-                    for k, v in storage_adapter.validate_json(f.read()).items()
+                    for k, v in tqdm(
+                        storage_adapter.validate_json(f.read()).items(),
+                        desc="Unpacking index",
+                    )
                 )
                 self._bm25_store = bm25
                 return bm25
@@ -115,11 +125,14 @@ class RAG:
                     question=question.question,
                     retrieved_sources=bm25.best_k(question.question, k),
                 )
-                for question in dataset.rag_questions
+                for question in tqdm(
+                    dataset.rag_questions, desc="Running querries"
+                )
             ],
             k=k,
         )
         try:
+            os.makedirs(save_directory, exist_ok=True)
             with open(
                 save_directory + "/" + os.path.basename(dataset_path), "w"
             ) as f:

diff --git a/src/rag/retrieval.py b/src/rag/retrieval.py
index 842a263c55e638360f7a62735337f0068a6bdc71..8e4dee36a18842d60756c37c078966290e90cc17 100644 (file)
--- a/src/rag/retrieval.py
+++ b/src/rag/retrieval.py
@@ -44,26 +44,29 @@ class BM25[T]:
 
     @staticmethod
     def from_storage(storage: Iterable[tuple[T, BagOfWords]]) -> "BM25[T]":
-        corpus = {k: (sum(b.values()), b) for k, b in storage}
+        corpus = {}
         word_usage: dict[str, set[T]] = {}
-        for k, (_, doc) in corpus.items():
+        corpus_words = 0
+        for k, doc in storage:
             for word in doc:
                 if word not in word_usage:
                     word_usage[word] = set()
                 word_usage[word].add(k)
-
-        return BM25(sum(doc[0] for doc in corpus.values()), corpus, word_usage)
+            doc_words = sum(doc.values())
+            corpus[k] = (doc_words, doc)
+            corpus_words += doc_words
+        return BM25(corpus_words, corpus, word_usage)
 
     def to_storage(self) -> Iterable[tuple[T, BagOfWords]]:
         return ((k, v[1]) for k, v in self.corpus.items())
 
     @staticmethod
     def from_corpus(
-        raw_corpus: dict[T, str], stopwords: set[str] | None = None
+        raw_corpus: Iterable[tuple[T, str]], stopwords: set[str] | None = None
     ) -> "BM25[T]":
         return BM25.from_storage(
             (k, bow)
-            for k, s in raw_corpus.items()
+            for k, s in raw_corpus
             if (bow := bag_of_words(s, stopwords if stopwords else set()))
         )
 
@@ -85,8 +88,30 @@ class BM25[T]:
     def score(self, querry: str, doc: T) -> float:
         return sum(self.word_score(q, doc) for q in words_normalize(querry))
 
+    def word_scores(self, word: str) -> dict[T, float]:
+        res = {}
+        k_plus_1 = self.k + 1
+        avgdl = self.corpus_words / len(self.corpus)
+        freq_bias_add = 1 - self.b
+        freq_bias_mul = self.b / avgdl
+        idf = self.idf(word)
+        for ident in self.word_usage.get(word, set()):
+            doc_words, doc = self.corpus[ident]
+            f_q = doc[word]
+            freq = f_q * k_plus_1
+            freq_bias = freq_bias_add + doc_words * freq_bias_mul
+            quotient = freq / (f_q + self.k * freq_bias)
+            res[ident] = idf * quotient
+        return res
+
     def scores(self, querry: str) -> dict[T, float]:
-        return {k: self.score(querry, k) for k in self.corpus}
+        # Old slow impl:
+        # return {k: self.score(querry, k) for k in self.corpus}
+        res = {}
+        for word in words_normalize(querry):
+            for k, v in self.word_scores(word).items():
+                res[k] = res.get(k, 0.0) + v
+        return res
 
     def best_k(self, querry: str, k: int) -> list[T]:
         k = min(len(self.corpus), k)