Merge pull request #876 from datalab-to/vik/quality

Vik/quality

Author: VikParuchuri

marker/converters/pdf.py

@@ -1,5 +1,7 @@
 import os
 
+from marker.schema.document import Document
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # disables a tokenizers warning
 
 from collections import defaultdict
@@ -171,7 +173,7 @@ def filepath_to_str(self, file_input: Union[str, io.BytesIO]):
             if temp_file is not None and os.path.exists(temp_file.name):
                 os.unlink(temp_file.name)
 
-    def build_document(self, filepath: str):
+    def build_document(self, filepath: str) -> Document:
         provider_cls = provider_from_filepath(filepath)
         layout_builder = self.resolve_dependencies(self.layout_builder_class)
         line_builder = self.resolve_dependencies(LineBuilder)

marker/models.py

@@ -1,5 +1,8 @@
 import os
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # Transformers uses .isin for an op, which is not supported on MPS
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = (
+    "1"  # Transformers uses .isin for an op, which is not supported on MPS
+)
 
 from surya.foundation import FoundationPredictor
 from surya.detection import DetectionPredictor
@@ -8,13 +11,18 @@
 from surya.recognition import RecognitionPredictor
 from surya.table_rec import TableRecPredictor
 
-def create_model_dict(device=None, dtype=None) -> dict:
-    foundation_predictor = FoundationPredictor(device=device, dtype=dtype)
+
+def create_model_dict(
+    device=None, dtype=None, attention_implementation: str | None = None
+) -> dict:
+    foundation_predictor = FoundationPredictor(
+        device=device, dtype=dtype, attention_implementation=attention_implementation
+    )
     return {
         "foundation_model": foundation_predictor,
         "layout_model": LayoutPredictor(device=device, dtype=dtype),
         "recognition_model": RecognitionPredictor(foundation_predictor),
         "table_rec_model": TableRecPredictor(device=device, dtype=dtype),
         "detection_model": DetectionPredictor(device=device, dtype=dtype),
-        "ocr_error_model": OCRErrorPredictor(device=device, dtype=dtype)
-    }
+        "ocr_error_model": OCRErrorPredictor(device=device, dtype=dtype),
+    }

marker/schema/blocks/base.py

@@ -100,6 +100,7 @@ class Block(BaseModel):
     lowres_image: Image.Image | None = None
     highres_image: Image.Image | None = None
     removed: bool = False  # Has block been replaced by new block?
+    _metadata: Optional[dict] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -114,6 +115,16 @@ def from_block(cls, block: Block) -> Block:
         block_attrs = block.model_dump(exclude=["id", "block_id", "block_type"])
         return cls(**block_attrs)
 
+    def set_internal_metadata(self, key, data):
+        if self._metadata is None:
+            self._metadata = {}
+        self._metadata[key] = data
+
+    def get_internal_metadata(self, key):
+        if self._metadata is None:
+            return None
+        return self._metadata.get(key)
+
     def get_image(
         self,
         document: Document,