update: colpali index syncs with wandb artifact

.gitignore

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

medrag_multi_modal/document_loader/load_image.py

@@ -3,11 +3,11 @@ import os
 from typing import Optional
 
 import rich
+import wandb
 import weave
 from pdf2image.pdf2image import convert_from_path
 from PIL import Image
 
-import wandb
 from medrag_multi_modal.document_loader.load_text import TextLoader
 
 

medrag_multi_modal/retrieval/__init__.py

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

medrag_multi_modal/retrieval/multi_modal_retrieval.py

@@ -1,22 +1,39 @@
+import os
+
+import wandb
 import weave
 from byaldi import RAGMultiModalModel
-import wandb
 
 
 class MultiModalRetriever(weave.Model):
     model_name: str
     _docs_retrieval_model: RAGMultiModalModel
-    
+
     def __init__(self, model_name: str = "vidore/colpali-v1.2"):
         super().__init__(model_name=model_name)
         self._docs_retrieval_model = RAGMultiModalModel.from_pretrained(self.model_name)
-    
+
     def index(self, data_artifact_name: str, weave_dataset_name: str, index_name: str):
         if wandb.run:
-            artifact = wandb.use_artifact(data_artifact_name, type='dataset')
+            artifact = wandb.use_artifact(data_artifact_name, type="dataset")
             artifact_dir = artifact.download()
         else:
             api = wandb.Api()
             artifact = api.artifact(data_artifact_name)
             artifact_dir = artifact.download()
-        self._docs_retrieval_model.index(input_path=artifact_dir, index_name=index_name)
+        self._docs_retrieval_model.index(
+            input_path=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()