add: BM25sRetriever

.gitignore

@@ -18,3 +18,5 @@ wandb/
 cursor_prompt.txt
 test.py
 uv.lock
+grays-anatomy-bm25s/
+prompt**.txt

medrag_multi_modal/retrieval/__init__.py

@@ -1,3 +1,4 @@
-from .multi_modal_retrieval import MultiModalRetriever
+from .bm25s_retrieval import BM25sRetriever
+from .colpali_retrieval import CalPaliRetriever
 
-__all__ = ["MultiModalRetriever"]
+__all__ = ["CalPaliRetriever", "BM25sRetriever"]

medrag_multi_modal/retrieval/bm25s_retrieval.py

@@ -0,0 +1,45 @@
+from typing import Optional
+
+import bm25s
+import weave
+from Stemmer import Stemmer
+
+import wandb
+
+LANGUAGE_DICT = {
+    "english": "en",
+    "french": "fr",
+    "german": "de",
+}
+
+
+class BM25sRetriever(weave.Model):
+    language: str
+    use_stemmer: bool
+    _retriever: Optional[bm25s.BM25]
+
+    def __init__(
+        self,
+        language: str = "english",
+        use_stemmer: bool = True,
+        retriever: Optional[bm25s.BM25] = None,
+    ):
+        super().__init__(language=language, use_stemmer=use_stemmer)
+        self._retriever = retriever or bm25s.BM25()
+
+    def index(self, corpus_dataset_name: str, index_name: Optional[str] = None):
+        corpus_dataset = weave.ref(corpus_dataset_name).get().rows
+        corpus = [row["text"] for row in corpus_dataset]
+        corpus_tokens = bm25s.tokenize(
+            corpus,
+            stopwords=LANGUAGE_DICT[self.language],
+            stemmer=Stemmer(self.language) if self.use_stemmer else None,
+        )
+        self._retriever.index(corpus_tokens)
+        self._retriever.save(index_name, corpus=[dict(row) for row in corpus_dataset])
+        if index_name:
+            self._retriever.save(index_name)
+            if wandb.run:
+                artifact = wandb.Artifact(name=index_name, type="bm25s-index")
+                artifact.add_dir(index_name)
+                artifact.save()

medrag_multi_modal/retrieval/colpali_retrieval.py

@@ -1,11 +1,12 @@
 import os
 from typing import Any, Optional
 
-import wandb
 import weave
 from byaldi import RAGMultiModalModel
 from PIL import Image
 
+import wandb
+
 from ..utils import get_wandb_artifact
 
 

pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "bm25s[full]>=0.2.2",
     "Byaldi>=0.0.5",
     "firerequests>=0.0.7",
+    "jax[cpu]>=0.4.34",
     "pdf2image>=1.17.0",
     "python-dotenv>=1.0.1",
     "pymupdf4llm>=0.0.17",
@@ -17,6 +18,7 @@ dependencies = [
     "uv>=0.4.20",
     "pytest>=8.3.3",
     "PyPDF2>=3.0.1",
+    "PyStemmer>=2.2.0.3",
     "isort>=5.13.2",
     "black>=24.10.0",
     "ruff>=0.6.9",
@@ -39,10 +41,12 @@ core = [
     "bm25s[full]>=0.2.2",
     "Byaldi>=0.0.5",
     "firerequests>=0.0.7",
+    "jax[cpu]>=0.4.34",
     "marker-pdf>=0.2.17",
     "pdf2image>=1.17.0",
     "pdfplumber>=0.11.4",
     "PyPDF2>=3.0.1",
+    "PyStemmer>=2.2.0.3",
     "python-dotenv>=1.0.1",
     "pymupdf4llm>=0.0.17",
     "semchunk>=2.2.0",