add: marker pdf text loader

medrag_multi_modal/__init__.py

@@ -1,5 +1,6 @@
 from .document_loader import (
     ImageLoader,
+    MarkerTextLoader,
     PDFPlumberTextLoader,
     PyMuPDF4LLMTextLoader,
     PyPDF2TextLoader,
@@ -11,6 +12,7 @@ __all__ = [
     "PyMuPDF4LLMTextLoader",
     "PyPDF2TextLoader",
     "PDFPlumberTextLoader",
+    "MarkerTextLoader",
     "ImageLoader",
     "TextImageLoader",
     "MultiModalRetriever",

medrag_multi_modal/document_loader/__init__.py

@@ -1,11 +1,17 @@
 from .load_image import ImageLoader
 from .load_text_image import TextImageLoader
-from .text_loader import PDFPlumberTextLoader, PyMuPDF4LLMTextLoader, PyPDF2TextLoader
+from .text_loader import (
+    MarkerTextLoader,
+    PDFPlumberTextLoader,
+    PyMuPDF4LLMTextLoader,
+    PyPDF2TextLoader,
+)
 
 __all__ = [
     "PyMuPDF4LLMTextLoader",
     "PyPDF2TextLoader",
     "PDFPlumberTextLoader",
+    "MarkerTextLoader",
     "ImageLoader",
     "TextImageLoader",
 ]

medrag_multi_modal/document_loader/text_loader/__init__.py

@@ -1,3 +1,4 @@
+from .marker_text_loader import MarkerTextLoader
 from .pdfplumber_text_loader import PDFPlumberTextLoader
 from .pymupdf4llm_text_loader import PyMuPDF4LLMTextLoader
 from .pypdf2_text_loader import PyPDF2TextLoader
@@ -6,4 +7,5 @@ __all__ = [
     "PyMuPDF4LLMTextLoader",
     "PyPDF2TextLoader",
     "PDFPlumberTextLoader",
+    "MarkerTextLoader",
 ]

medrag_multi_modal/document_loader/text_loader/marker_text_loader.py

@@ -0,0 +1,29 @@
+from typing import Dict
+
+from marker.convert import convert_single_pdf
+from marker.models import load_all_models
+
+from .base_text_loader import BaseTextLoader
+
+
+class MarkerTextLoader(BaseTextLoader):
+    async def _process_page(self, page_idx: int) -> Dict[str, str]:
+        model_lst = load_all_models()
+
+        text, _, out_meta = convert_single_pdf(
+            self.document_file_path,
+            model_lst,
+            max_pages=1,
+            batch_multiplier=1,
+            start_page=page_idx,
+            ocr_all_pages=True,
+        )
+
+        return {
+            "text": text,
+            "meta": out_meta,
+            "page_idx": page_idx,
+            "document_name": self.document_name,
+            "file_path": self.document_file_path,
+            "file_url": self.url,
+        }