Align with surya refactor

marker/builders/layout.py

@@ -1,11 +1,10 @@
 from typing import Annotated, List, Optional, Tuple
 
 import numpy as np
-from surya.layout import batch_layout_detection
-from surya.model.layout.encoderdecoder import SuryaLayoutModel
-from surya.model.ocr_error.model import DistilBertForSequenceClassification
-from surya.ocr_error import batch_ocr_error_detection
-from surya.schema import LayoutResult, OCRErrorDetectionResult
+from surya.layout import LayoutPredictor
+from surya.layout.schema import LayoutResult
+from surya.ocr_error import OCRErrorPredictor
+from surya.ocr_error.schema import OCRErrorDetectionResult
 
 from marker.builders import BaseBuilder
 from marker.providers import ProviderOutput, ProviderPageLines
@@ -52,7 +51,7 @@ class LayoutBuilder(BaseBuilder):
         "A list of block types to exclude from the layout coverage check.",
     ] = (BlockTypes.Figure, BlockTypes.Picture, BlockTypes.Table, BlockTypes.FigureGroup, BlockTypes.TableGroup, BlockTypes.PictureGroup)
 
-    def __init__(self, layout_model: SuryaLayoutModel, ocr_error_model: DistilBertForSequenceClassification, config=None):
+    def __init__(self, layout_model: LayoutPredictor, ocr_error_model: OCRErrorPredictor, config=None):
         self.layout_model = layout_model
         self.ocr_error_model = ocr_error_model
 
@@ -71,11 +70,8 @@ class LayoutBuilder(BaseBuilder):
         return 6
 
     def surya_layout(self, pages: List[PageGroup]) -> List[LayoutResult]:
-        processor = self.layout_model.processor
-        layout_results = batch_layout_detection(
+        layout_results = self.layout_model(
             [p.lowres_image for p in pages],
-            self.layout_model,
-            processor,
             batch_size=int(self.get_batch_size())
         )
         return layout_results
@@ -97,10 +93,8 @@ class LayoutBuilder(BaseBuilder):
 
             page_texts.append(page_text)
 
-        ocr_error_detection_results = batch_ocr_error_detection(
+        ocr_error_detection_results = self.ocr_error_model(
             page_texts,
-            self.ocr_error_model,
-            self.ocr_error_model.tokenizer,
             batch_size=int(self.get_batch_size())  # TODO Better Multiplier
         )
         return ocr_error_detection_results

marker/builders/llm_layout.py

@@ -3,8 +3,8 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Annotated, Optional
 
 from google.ai.generativelanguage_v1beta.types import content
-from surya.model.layout.encoderdecoder import SuryaLayoutModel
-from surya.model.ocr_error.model import DistilBertForSequenceClassification
+from surya.layout import LayoutPredictor
+from surya.ocr_error import OCRErrorPredictor
 from tqdm import tqdm
 
 from marker.builders.layout import LayoutBuilder
@@ -91,7 +91,7 @@ Respond only with one of `Figure`, `Picture`, `ComplexRegion`, `Table`, or `Form
 Here is the image of the layout block:
 """
 
-    def __init__(self, layout_model: SuryaLayoutModel, ocr_error_model: DistilBertForSequenceClassification, config=None):
+    def __init__(self, layout_model: LayoutPredictor, ocr_error_model: OCRErrorPredictor, config=None):
         super().__init__(layout_model, ocr_error_model, config)
 
         self.model = GoogleModel(self.google_api_key, self.model_name)

marker/builders/ocr.py

@@ -1,9 +1,8 @@
 from typing import Annotated, List, Optional
 
 from ftfy import fix_text
-from surya.model.detection.model import EfficientViTForSemanticSegmentation
-from surya.model.recognition.encoderdecoder import OCREncoderDecoderModel
-from surya.ocr import run_ocr
+from surya.detection import DetectionPredictor
+from surya.recognition import RecognitionPredictor
 
 from marker.builders import BaseBuilder
 from marker.providers import ProviderOutput, ProviderPageLines
@@ -37,7 +36,7 @@ class OcrBuilder(BaseBuilder):
         "Default is None."
     ] = None
 
-    def __init__(self, detection_model: EfficientViTForSemanticSegmentation, recognition_model: OCREncoderDecoderModel, config=None):
+    def __init__(self, detection_model: DetectionPredictor, recognition_model: RecognitionPredictor, config=None):
         super().__init__(config)
 
         self.detection_model = detection_model
@@ -65,13 +64,10 @@ class OcrBuilder(BaseBuilder):
 
     def ocr_extraction(self, document: Document, provider: PdfProvider) -> ProviderPageLines:
         page_list = [page for page in document.pages if page.text_extraction_method == "surya"]
-        recognition_results = run_ocr(
+        recognition_results = self.recognition_model(
             images=[page.lowres_image for page in page_list],
             langs=[self.languages] * len(page_list),
-            det_model=self.detection_model,
-            det_processor=self.detection_model.processor,
-            rec_model=self.recognition_model,
-            rec_processor=self.recognition_model.processor,
+            det_predictor=self.detection_model,
             detection_batch_size=int(self.get_detection_batch_size()),
             recognition_batch_size=int(self.get_recognition_batch_size()),
             highres_images=[page.highres_image for page in page_list]

marker/models.py

@@ -1,86 +1,49 @@
 import os
 
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # Transformers uses .isin for a simple op, which is not supported on MPS
-
-from surya.model.detection.model import load_model as load_detection_model, load_processor as load_detection_processor
-from surya.model.layout.model import load_model as load_layout_model
-from surya.model.layout.processor import load_processor as load_layout_processor
-from texify.model.model import load_model as load_texify_model
-from texify.model.processor import load_processor as load_texify_processor
 from marker.settings import settings
-from surya.model.recognition.model import load_model as load_recognition_model
-from surya.model.recognition.processor import load_processor as load_recognition_processor
-from surya.model.table_rec.model import load_model as load_table_model
-from surya.model.table_rec.processor import load_processor as load_table_processor
-from surya.model.ocr_error.model import load_model as load_ocr_error_model
-from surya.model.ocr_error.model import load_tokenizer as load_ocr_error_tokenizer
-
-from texify.model.model import GenerateVisionEncoderDecoderModel
-from surya.model.layout.encoderdecoder import SuryaLayoutModel
-from surya.model.detection.model import EfficientViTForSemanticSegmentation
-from surya.model.recognition.encoderdecoder import OCREncoderDecoderModel
-from surya.model.table_rec.encoderdecoder import TableRecEncoderDecoderModel
-from surya.model.ocr_error.model import DistilBertForSequenceClassification
-
 
-def setup_table_rec_model(device=None, dtype=None) -> TableRecEncoderDecoderModel:
-    if device:
-        table_model = load_table_model(device=device, dtype=dtype)
-    else:
-        table_model = load_table_model()
-    table_model.processor = load_table_processor()
-    return table_model
-
-
-def setup_recognition_model(device=None, dtype=None) -> OCREncoderDecoderModel:
-    if device:
-        rec_model = load_recognition_model(device=device, dtype=dtype)
-    else:
-        rec_model = load_recognition_model()
-    rec_model.processor = load_recognition_processor()
-    return rec_model
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # Transformers uses .isin for a simple op, which is not supported on MPS
 
+from typing import List
+from PIL import Image
 
-def setup_detection_model(device=None, dtype=None) -> EfficientViTForSemanticSegmentation:
-    if device:
-        model = load_detection_model(device=device, dtype=dtype)
-    else:
-        model = load_detection_model()
-    model.processor = load_detection_processor()
-    return model
+from surya.detection import DetectionPredictor
+from surya.layout import LayoutPredictor
+from surya.ocr_error import OCRErrorPredictor
+from surya.recognition import RecognitionPredictor
+from surya.table_rec import TableRecPredictor
 
+from texify.model.model import load_model as load_texify_model
+from texify.model.processor import load_processor as load_texify_processor
+from texify.inference import batch_inference
 
-def setup_texify_model(device=None, dtype=None) -> GenerateVisionEncoderDecoderModel:
-    if device:
-        texify_model = load_texify_model(checkpoint=settings.TEXIFY_MODEL_NAME, device=device, dtype=dtype)
-    else:
-        texify_model = load_texify_model(checkpoint=settings.TEXIFY_MODEL_NAME, device=settings.TORCH_DEVICE_MODEL, dtype=settings.TEXIFY_DTYPE)
-    texify_model.processor = load_texify_processor()
-    return texify_model
+class TexifyPredictor:
+    def __init__(self, device=None, dtype=None):
+        if not device:
+            device = settings.TORCH_DEVICE_MODEL
+        if not dtype:
+            dtype = settings.TEXIFY_DTYPE
 
+        self.model = load_texify_model(checkpoint=settings.TEXIFY_MODEL_NAME, device=device, dtype=dtype)
+        self.processor = load_texify_processor()
+        self.device = device
+        self.dtype = dtype
 
-def setup_layout_model(device=None, dtype=None) -> SuryaLayoutModel:
-    if device:
-        model = load_layout_model(device=device, dtype=dtype)
-    else:
-        model = load_layout_model()
-    model.processor = load_layout_processor()
-    return model
+    def __call__(self, batch_images: List[Image.Image], max_tokens: int):
+        return batch_inference(
+            batch_images,
+            self.model,
+            self.processor,
+            max_tokens=max_tokens
+        )
 
-def setup_ocr_error_model(device=None, dtype=None) -> DistilBertForSequenceClassification:
-    if device:
-        model = load_ocr_error_model(device=device, dtype=dtype)
-    else:
-        model = load_ocr_error_model()
-    model.tokenizer = load_ocr_error_tokenizer()
-    return model
 
 def create_model_dict(device=None, dtype=None) -> dict:
     return {
-        "layout_model": setup_layout_model(device, dtype),
-        "texify_model": setup_texify_model(device, dtype),
-        "recognition_model": setup_recognition_model(device, dtype),
-        "table_rec_model": setup_table_rec_model(device, dtype),
-        "detection_model": setup_detection_model(device, dtype),
-        "ocr_error_model": setup_ocr_error_model(device,dtype)
+        "layout_model": LayoutPredictor(device=device, dtype=dtype),
+        "texify_model": TexifyPredictor(device=device, dtype=dtype),
+        "recognition_model": RecognitionPredictor(device=device, dtype=dtype),
+        "table_rec_model": TableRecPredictor(device=device, dtype=dtype),
+        "detection_model": DetectionPredictor(device=device, dtype=dtype),
+        "ocr_error_model": OCRErrorPredictor(device=device, dtype=dtype)
     }

marker/processors/equation.py

@@ -4,6 +4,7 @@ from texify.inference import batch_inference
 from texify.model.model import GenerateVisionEncoderDecoderModel
 from tqdm import tqdm
 
+from marker.models import TexifyPredictor
 from marker.processors import BaseProcessor
 from marker.schema import BlockTypes
 from marker.schema.document import Document
@@ -32,7 +33,7 @@ class EquationProcessor(BaseProcessor):
         "The number of tokens to buffer above max for the Texify model.",
     ] = 256
 
-    def __init__(self, texify_model: GenerateVisionEncoderDecoderModel, config=None):
+    def __init__(self, texify_model: TexifyPredictor, config=None):
         super().__init__(config)
 
         self.texify_model = texify_model
@@ -92,10 +93,8 @@ class EquationProcessor(BaseProcessor):
 
             batch_images = [eq["image"] for eq in batch_equations]
 
-            model_output = batch_inference(
+            model_output = self.texify_model(
                 batch_images,
-                self.texify_model,
-                self.texify_model.processor,
                 max_tokens=max_length
             )
 

marker/processors/table.py

@@ -2,10 +2,9 @@
 from typing import Annotated
 
 from ftfy import fix_text
-from surya.input.pdflines import get_page_text_lines
-from surya.model.detection.model import EfficientViTForSemanticSegmentation
-from surya.model.recognition.encoderdecoder import OCREncoderDecoderModel
-from surya.model.table_rec.encoderdecoder import TableRecEncoderDecoderModel
+from surya.detection import DetectionPredictor
+from surya.recognition import RecognitionPredictor
+from surya.table_rec import TableRecPredictor
 from tabled.assignment import assign_rows_columns
 from tabled.inference.recognition import get_cells, recognize_tables
 
@@ -42,9 +41,9 @@ class TableProcessor(BaseProcessor):
 
     def __init__(
         self,
-        detection_model: EfficientViTForSemanticSegmentation,
-        recognition_model: OCREncoderDecoderModel,
-        table_rec_model: TableRecEncoderDecoderModel,
+        detection_model: DetectionPredictor,
+        recognition_model: RecognitionPredictor,
+        table_rec_model: TableRecPredictor,
         config=None
     ):
         super().__init__(config)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "marker-pdf"
-version = "1.2.3"
+version = "1.3.0"
 description = "Convert PDF to markdown with high speed and accuracy."
 authors = ["Vik Paruchuri <github@vikas.sh>"]
 readme = "README.md"

tests/builders/test_blank_page.py

@@ -1,4 +1,4 @@
-from surya.schema import LayoutResult
+from surya.layout.schema import LayoutResult
 
 from marker.builders.document import DocumentBuilder
 from marker.builders.layout import LayoutBuilder

tests/conftest.py

@@ -9,9 +9,7 @@ from marker.builders.document import DocumentBuilder
 from marker.builders.layout import LayoutBuilder
 from marker.builders.ocr import OcrBuilder
 from marker.converters.pdf import PdfConverter
-from marker.models import setup_detection_model, setup_layout_model, \
-    setup_recognition_model, setup_table_rec_model, \
-    setup_texify_model, setup_ocr_error_model
+from marker.models import create_model_dict
 from marker.schema import BlockTypes
 from marker.schema.blocks import Block
 from marker.renderers.markdown import MarkdownRenderer
@@ -19,46 +17,42 @@ from marker.renderers.json import JSONRenderer
 from marker.schema.registry import register_block_class
 from marker.util import classes_to_strings
 
+@pytest.fixture(scope="session")
+def model_dict():
+    model_dict = create_model_dict()
+    yield model_dict
+    del model_dict
+
 
 @pytest.fixture(scope="session")
-def layout_model():
-    layout_m = setup_layout_model()
-    yield layout_m
-    del layout_m
+def layout_model(model_dict):
+    yield model_dict["layout_model"]
 
 
 @pytest.fixture(scope="session")
-def detection_model():
-    detection_m = setup_detection_model()
-    yield detection_m
-    del detection_m
+def detection_model(model_dict):
+    yield model_dict["detection_model"]
 
 
 @pytest.fixture(scope="session")
-def texify_model():
-    texify_m = setup_texify_model()
-    yield texify_m
-    del texify_m
+def texify_model(model_dict):
+    yield model_dict["texify_model"]
 
 
 @pytest.fixture(scope="session")
-def recognition_model():
-    ocr_m = setup_recognition_model()
-    yield ocr_m
-    del ocr_m
+def recognition_model(model_dict):
+    yield model_dict["recognition_model"]
 
 
 @pytest.fixture(scope="session")
-def table_rec_model():
-    table_rec_m = setup_table_rec_model()
-    yield table_rec_m
-    del table_rec_m
+def table_rec_model(model_dict):
+    yield model_dict["table_rec_model"]
+
 
 @pytest.fixture(scope="session")
-def ocr_error_model():
-    ocr_error_m = setup_ocr_error_model()
-    yield ocr_error_m
-    del ocr_error_m
+def ocr_error_model(model_dict):
+    yield model_dict["ocr_error_model"]
+
 
 @pytest.fixture(scope="function")
 def config(request):
@@ -101,15 +95,7 @@ def pdf_document(request, config, pdf_provider, layout_model, ocr_error_model, r
 
 
 @pytest.fixture(scope="function")
-def pdf_converter(request, config, layout_model, texify_model, recognition_model, table_rec_model, detection_model, ocr_error_model, renderer):
-    model_dict = {
-        "layout_model": layout_model,
-        "texify_model": texify_model,
-        "recognition_model": recognition_model,
-        "table_rec_model": table_rec_model,
-        "detection_model": detection_model,
-        "ocr_error_model": ocr_error_model
-    }
+def pdf_converter(request, config, model_dict, renderer):
     yield PdfConverter(
         artifact_dict=model_dict,
         processor_list=None,

tests/utils.py

@@ -2,11 +2,6 @@ from marker.providers.pdf import PdfProvider
 import tempfile
 
 import datasets
-from marker.models import setup_layout_model, setup_recognition_model, setup_detection_model
-from marker.builders.document import DocumentBuilder
-from marker.builders.layout import LayoutBuilder
-from marker.builders.ocr import OcrBuilder
-from marker.schema.document import Document
 
 
 def setup_pdf_provider(