integrate changes and increment surya version

convert_single.py

@@ -1,6 +1,6 @@
 import os
 
-os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # Transformers uses .isin for a simple op, which is not supported on MPS
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # Transformers uses .isin for a simple op, which is not supported on MPS
 
 import time
 

marker/builders/high_quality_layout.py

@@ -57,16 +57,16 @@ Here are the top k predictions from the model followed by the image:
     def __init__(self, layout_model: SuryaLayoutModel, config=None):
         super().__init__(layout_model, config)
 
-        if self.google_api_key is not None:
-            genai.configure(api_key=self.google_api_key)
-            self.model = genai.GenerativeModel(self.model_name)
+        self.model = None
+        if self.google_api_key is None:
+            raise ValueError("Google API key is not set")
+        
+        genai.configure(api_key=self.google_api_key)
+        self.model = genai.GenerativeModel(self.model_name)
 
     def __call__(self, document: Document, provider: PdfProvider):
         super().__call__(document, provider)
 
-        if self.model is None:
-            return
-
         self.relabel_blocks(document)
 
     def relabel_blocks(self, document: Document):

marker/config/parser.py

@@ -34,6 +34,7 @@ class ConfigParser:
         fn = click.option("--disable_multiprocessing", is_flag=True, default=False, help="Disable multiprocessing.")(fn)
         fn = click.option("--paginate_output", is_flag=True, default=False, help="Paginate output.")(fn)
         fn = click.option("--disable_image_extraction", is_flag=True, default=False, help="Disable image extraction.")(fn)
+        fn = click.option("--high_quality", is_flag=True, default=False, help="Enable high quality processing with Gemini.")(fn)
         return fn
 
     def generate_config_dict(self) -> Dict[str, any]:
@@ -69,6 +70,9 @@ class ConfigParser:
                 case "disable_image_extraction":
                     if v:
                         config["extract_images"] = False
+                case "high_quality":
+                    if v:
+                        config["high_quality"] = True
         return config
 
     def get_renderer(self):

marker/converters/pdf.py

@@ -45,6 +45,7 @@ class PdfConverter(BaseConverter):
             instead of the defaults.
     """
     override_map: Dict[BlockTypes, Type[Block]] = defaultdict()
+    high_quality: bool = False
 
     def __init__(self, artifact_dict: Dict[str, Any], processor_list: List[str] | None = None, renderer: str | None = None, config=None):
         super().__init__(config)
@@ -81,6 +82,10 @@ class PdfConverter(BaseConverter):
         self.processor_list = processor_list
         self.renderer = renderer
 
+        self.layout_builder_class = LayoutBuilder
+        if self.high_quality:
+            self.layout_builder_class = HighQualityLayoutBuilder
+
     def resolve_dependencies(self, cls):
         init_signature = inspect.signature(cls.__init__)
         parameters = init_signature.parameters
@@ -102,7 +107,7 @@ class PdfConverter(BaseConverter):
 
     def __call__(self, filepath: str):
         pdf_provider = PdfProvider(filepath, self.config)
-        layout_builder = self.resolve_dependencies(HighQualityLayoutBuilder)
+        layout_builder = self.resolve_dependencies(self.layout_builder_class)
         ocr_builder = self.resolve_dependencies(OcrBuilder)
         document = DocumentBuilder(self.config)(pdf_provider, layout_builder, ocr_builder)
         StructureBuilder(self.config)(document)

marker/processors/high_quality_text.py

@@ -40,6 +40,7 @@ class HighQualityTextProcessor(BaseProcessor):
     google_api_key: Optional[str] = settings.GOOGLE_API_KEY
     confidence_threshold: float = 0.7
     model_name: str = "gemini-1.5-flash"
+    high_quality: bool = False
 
     gemini_rewriting_prompt = """You are a text correction expert specializing in accurately reproducing text from images.
 You will receive an image of a text block and a set of extracted lines corresponding to the text in the image.
@@ -98,12 +99,18 @@ Output:
     def __init__(self, config=None):
         super().__init__(config)
 
-        if self.google_api_key is not None:
-            genai.configure(api_key=self.google_api_key)
-            self.model = genai.GenerativeModel(self.model_name)
+        self.model = None
+        if not self.high_quality:
+            return
+
+        if self.google_api_key is None:
+            raise ValueError("Google API key is not set")
+
+        genai.configure(api_key=self.google_api_key)
+        self.model = genai.GenerativeModel(self.model_name)
 
     def __call__(self, document: Document):
-        if self.model is None:
+        if not self.high_quality or self.model is None:
             return
 
         self.rewrite_blocks(document)

pyproject.toml

@@ -33,7 +33,7 @@ tabulate = "^0.9.0"
 ftfy = "^6.1.1"
 texify = "^0.2.1"
 rapidfuzz = "^3.8.1"
-surya-ocr = "~0.8.0"
+surya-ocr = "~0.8.1"
 regex = "^2024.4.28"
 pdftext = "~0.4.0"
 tabled-pdf = "~0.2.0"