Merge pull request #4 from soumik12345/feat/colpali-retrieval

.gitignore

@@ -7,4 +7,6 @@ cursor_prompt.txt
 test.py
 **.pdf
 images/
-wandb/
+wandb/
+.byaldi/
+artifacts/

README.md

@@ -1,3 +1,3 @@
 # MedRAG Multi-Modal
 
-Multi-modal RAG for medical docmain.
+Multi-modal RAG for medical docmain.

docs/index.md

@@ -1,3 +1,40 @@
 # MedRAG Multi-Modal
 
-Multi-modal RAG for medical docmain.
+Multi-modal RAG for medical docmain.
+
+## Installation
+
+### For Development
+
+For MacOS, you need to run
+
+```bash
+brew install poppler
+```
+
+For Debian/Ubuntu, you need to run
+
+```bash
+sudo apt-get install -y poppler-utils
+```
+
+Then, you can install the dependencies using uv in the virtual environment `.venv` using
+
+```bash
+git clone https://github.com/soumik12345/medrag-multi-modal
+cd medrag-multi-modal
+pip install -U pip uv
+uv sync
+```
+
+After this, you need to activate the virtual environment using
+
+```bash
+source .venv/bin/activate
+```
+
+In the activated virtual environment, you can optionally install Flash Attention (required for ColPali) using
+
+```bash
+uv pip install flash-attn --no-build-isolation
+```

docs/installation/development.md

@@ -0,0 +1,40 @@
+# Setting up the development environment
+
+## Install Poppler
+
+For MacOS, you need to run
+
+```bash
+brew install poppler
+```
+
+For Debian/Ubuntu, you need to run
+
+```bash
+sudo apt-get install -y poppler-utils
+```
+
+## Install the dependencies
+
+Then, you can install the dependencies using uv in the virtual environment `.venv` using
+
+```bash
+git clone https://github.com/soumik12345/medrag-multi-modal
+cd medrag-multi-modal
+pip install -U pip uv
+uv sync
+```
+
+After this, you need to activate the virtual environment using
+
+```bash
+source .venv/bin/activate
+```
+
+## [Optional] Install Flash Attention
+
+In the activated virtual environment, you can optionally install Flash Attention (required for ColPali) using
+
+```bash
+uv pip install flash-attn --no-build-isolation
+```

docs/installation/install.md

@@ -0,0 +1,9 @@
+# Installation
+
+You just need to clone the repository and run the install.sh script
+
+```bash
+git clone https://github.com/soumik12345/medrag-multi-modal
+cd medrag-multi-modal
+sh install.sh
+```

docs/retreival/multi_modal_retrieval.md

@@ -0,0 +1,3 @@
+# Multi-Modal Retrieval
+
+::: medrag_multi_modal.retrieval.multi_modal_retrieval

install.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+OS_TYPE=$(uname -s)
+
+if [ "$OS_TYPE" = "Darwin" ]; then
+    echo "Detected macOS."
+    brew install poppler
+elif [ "$OS_TYPE" = "Linux" ]; then
+    if [ -f /etc/os-release ]; then
+        . /etc/os-release
+        if [ "$ID" = "ubuntu" ] || [ "$ID" = "debian" ]; then
+            echo "Detected Ubuntu/Debian."
+            sudo apt-get update
+            sudo apt-get install -y poppler-utils
+        else
+            echo "Unsupported Linux distribution: $ID"
+            exit 1
+        fi
+    else
+        echo "Cannot detect Linux distribution."
+        exit 1
+    fi
+else
+    echo "Unsupported OS: $OS_TYPE"
+    exit 1
+fi
+
+git clone https://github.com/soumik12345/medrag-multi-modal
+cd medrag-multi-modal
+pip install -U .[core]

medrag_multi_modal/document_loader/load_image.py

@@ -13,8 +13,8 @@ from medrag_multi_modal.document_loader.load_text import TextLoader
 
 class ImageLoader(TextLoader):
     """
-    ImageLoader is a class that extends the `TextLoader` class to handle the extraction and
-    loading of images from a PDF file.
+    `ImageLoader` is a class that extends the `TextLoader` class to handle the extraction and
+    loading of pages from a PDF file as images.
 
     This class provides functionality to convert specific pages of a PDF document into images
     and optionally publish these images to a Weave dataset.

medrag_multi_modal/retrieval/__init__.py

@@ -0,0 +1,3 @@
+from .multi_modal_retrieval import MultiModalRetriever
+
+__all__ = ["MultiModalRetriever"]

medrag_multi_modal/retrieval/multi_modal_retrieval.py

@@ -0,0 +1,149 @@
+import os
+from typing import Any, Optional
+
+import weave
+from byaldi import RAGMultiModalModel
+from PIL import Image
+
+import wandb
+
+from ..utils import get_wandb_artifact
+
+
+class MultiModalRetriever(weave.Model):
+    """
+    MultiModalRetriever is a class that facilitates the retrieval of page images using ColPali.
+    
+    This class leverages the `byaldi.RAGMultiModalModel` to perform document retrieval tasks.
+    It can be initialized with a pre-trained model or from a specified W&B artifact. The class
+    also provides methods to index new data and to predict/retrieve documents based on a query.
+    
+    !!! example "Indexing Data"
+        ```python
+        import wandb
+        from medrag_multi_modal.retrieval import MultiModalRetriever
+
+        wandb.init(project="medrag-multi-modal", entity="ml-colabs", job_type="index")
+        retriever = MultiModalRetriever()
+        retriever.index(
+            data_artifact_name="ml-colabs/medrag-multi-modal/grays-anatomy-images:v1",
+            weave_dataset_name="grays-anatomy-images:v0",
+            index_name="grays-anatomy",
+        )
+        ```
+    
+    !!! example "Retrieving Documents"
+        ```python
+        import weave
+
+        import wandb
+        from medrag_multi_modal.retrieval import MultiModalRetriever
+        
+        weave.init(project_name="ml-colabs/medrag-multi-modal")
+        retriever = MultiModalRetriever.from_artifact(
+            index_artifact_name="ml-colabs/medrag-multi-modal/grays-anatomy:v0",
+            metadata_dataset_name="grays-anatomy-images:v0",
+            data_artifact_name="ml-colabs/medrag-multi-modal/grays-anatomy-images:v1",
+        )
+        retriever.predict(
+            query="which neurotransmitters convey information between Merkel cells and sensory afferents?",
+            top_k=3,
+        )
+        ```
+
+    Attributes:
+        model_name (str): The name of the model to be used for retrieval.
+    """
+    model_name: str
+    _docs_retrieval_model: Optional[RAGMultiModalModel] = None
+    _metadata: Optional[dict] = None
+    _data_artifact_dir: Optional[str] = None
+
+    def __init__(
+        self,
+        model_name: str = "vidore/colpali-v1.2",
+        docs_retrieval_model: Optional[RAGMultiModalModel] = None,
+        data_artifact_dir: Optional[str] = None,
+        metadata_dataset_name: Optional[str] = None,
+    ):
+        super().__init__(model_name=model_name)
+        self._docs_retrieval_model = (
+            docs_retrieval_model or RAGMultiModalModel.from_pretrained(self.model_name)
+        )
+        self._data_artifact_dir = data_artifact_dir
+        self._metadata = (
+            [dict(row) for row in weave.ref(metadata_dataset_name).get().rows]
+            if metadata_dataset_name
+            else None
+        )
+
+    @classmethod
+    def from_artifact(
+        cls,
+        index_artifact_name: str,
+        metadata_dataset_name: str,
+        data_artifact_name: str,
+    ):
+        index_artifact_dir = get_wandb_artifact(index_artifact_name, "colpali-index")
+        data_artifact_dir = get_wandb_artifact(data_artifact_name, "dataset")
+        docs_retrieval_model = RAGMultiModalModel.from_index(
+            index_path=os.path.join(index_artifact_dir, "index")
+        )
+        return cls(
+            docs_retrieval_model=docs_retrieval_model,
+            metadata_dataset_name=metadata_dataset_name,
+            data_artifact_dir=data_artifact_dir,
+        )
+
+    def index(self, data_artifact_name: str, weave_dataset_name: str, index_name: str):
+        data_artifact_dir = get_wandb_artifact(data_artifact_name, "dataset")
+        self._docs_retrieval_model.index(
+            input_path=data_artifact_dir,
+            index_name=index_name,
+            store_collection_with_index=False,
+            overwrite=True,
+        )
+        if wandb.run:
+            artifact = wandb.Artifact(
+                name=index_name,
+                type="colpali-index",
+                metadata={"weave_dataset_name": weave_dataset_name},
+            )
+            artifact.add_dir(
+                local_path=os.path.join(".byaldi", index_name), name="index"
+            )
+            artifact.save()
+
+    @weave.op()
+    def predict(self, query: str, top_k: int = 3) -> list[dict[str, Any]]:
+        """
+        Predicts and retrieves the top-k most relevant documents/images for a given query
+        using ColPali.
+
+        This function uses the document retrieval model to search for the most relevant
+        documents based on the provided query. It returns a list of dictionaries, each
+        containing the document image, document ID, and the relevance score.
+
+        Args:
+            query (str): The search query string.
+            top_k (int, optional): The number of top results to retrieve. Defaults to 10.
+
+        Returns:
+            list[dict[str, Any]]: A list of dictionaries where each dictionary contains:
+                - "doc_image" (PIL.Image.Image): The image of the document.
+                - "doc_id" (str): The ID of the document.
+                - "score" (float): The relevance score of the document.
+        """
+        results = self._docs_retrieval_model.search(query=query, k=top_k)
+        retrieved_results = []
+        for result in results:
+            retrieved_results.append(
+                {
+                    "doc_image": Image.open(
+                        os.path.join(self._data_artifact_dir, f"{result['doc_id']}.png")
+                    ),
+                    "doc_id": result["doc_id"],
+                    "score": result["score"],
+                }
+            )
+        return retrieved_results

medrag_multi_modal/utils.py

@@ -0,0 +1,12 @@
+import wandb
+
+
+def get_wandb_artifact(artifact_name: str, artifact_type: str) -> str:
+    if wandb.run:
+        artifact = wandb.use_artifact(artifact_name, type=artifact_type)
+        artifact_dir = artifact.download()
+    else:
+        api = wandb.Api()
+        artifact = api.artifact(artifact_name)
+        artifact_dir = artifact.download()
+    return artifact_dir

mkdocs.yml

@@ -59,9 +59,14 @@ extra_javascript:
 
 nav:
   - Home: 'index.md'
+  - Setup:
+    - Installation: 'installation/install.md'
+    - Development: 'installation/development.md'
   - Document Loader:
     - Text Loader: 'document_loader/load_text.md'
     - Text and Image Loader: 'document_loader/load_text_image.md'
     - Image Loader: 'document_loader/load_image.md'
+  - Retrieval:
+    - Multi-Modal Retrieval: 'retreival/multi_modal_retrieval.md'
 
-repo_url: https://github.com/soumik12345/medrag-multi-modal
+repo_url: https://github.com/soumik12345/medrag-multi-modal