update: BaseImageLoader + MarkerImageLoader

medrag_multi_modal/document_loader/image_loader/base_img_loader.py

@@ -3,6 +3,7 @@ import os
 from abc import abstractmethod
 from typing import Dict, List, Optional
 
+import jsonlines
 import rich
 
 import wandb
@@ -41,7 +42,8 @@ class BaseImageLoader(BaseTextLoader):
         end_page: Optional[int] = None,
         wandb_artifact_name: Optional[str] = None,
         image_save_dir: str = "./images",
-        cleanup: bool = True,
+        exclude_file_extensions: list[str] = [],
+        cleanup: bool = False,
         **kwargs,
     ) -> List[Dict[str, str]]:
         """
@@ -61,10 +63,11 @@ class BaseImageLoader(BaseTextLoader):
         If a wandb_artifact_name is provided, the processed pages are published to a WandB artifact.
 
         Args:
-            start_page (Optional[int]): The starting page index (0-based) to process. Defaults to the first page.
-            end_page (Optional[int]): The ending page index (0-based) to process. Defaults to the last page.
+            start_page (Optional[int]): The starting page index (0-based) to process.
+            end_page (Optional[int]): The ending page index (0-based) to process.
             wandb_artifact_name (Optional[str]): The name of the WandB artifact to publish the pages to, if provided.
             image_save_dir (str): The directory to save the extracted images.
+            exclude_file_extensions (list[str]): A list of file extensions to exclude from the image_save_dir.
             cleanup (bool): Whether to remove extracted images from `image_save_dir`, if uploading to wandb artifact.
             **kwargs: Additional keyword arguments that will be passed to extract_page_data method and the underlying library.
 
@@ -99,6 +102,15 @@ class BaseImageLoader(BaseTextLoader):
         for task in asyncio.as_completed(tasks):
             await task
 
+        with jsonlines.open(
+            os.path.join(image_save_dir, "metadata.jsonl"), mode="w"
+        ) as writer:
+            writer.write(pages)
+
+        for file in os.listdir(image_save_dir):
+            if file.endswith(tuple(exclude_file_extensions)):
+                os.remove(os.path.join(image_save_dir, file))
+
         if wandb_artifact_name:
             artifact = wandb.Artifact(
                 name=wandb_artifact_name,

medrag_multi_modal/document_loader/image_loader/marker_img_loader.py

@@ -97,6 +97,6 @@ class MarkerImageLoader(BaseImageLoader):
             "document_name": self.document_name,
             "file_path": self.document_file_path,
             "file_url": self.url,
-            "image_file_paths": image_file_paths,
+            "image_file_paths": os.path.join(image_save_dir, "*.png"),
             "meta": out_meta,
         }

pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
     "google-generativeai>=0.8.3",
     "mistralai>=1.1.0",
     "instructor>=1.6.3",
+    "jsonlines>=4.0.0",
 ]
 
 [project.optional-dependencies]
@@ -67,6 +68,7 @@ core = [
     "google-generativeai>=0.8.3",
     "mistralai>=1.1.0",
     "instructor>=1.6.3",
+    "jsonlines>=4.0.0",
 ]
 
 dev = ["pytest>=8.3.3", "isort>=5.13.2", "black>=24.10.0", "ruff>=0.6.9"]