Merge pull request #229 from VikParuchuri/dev

convert_single.py

@@ -1,3 +1,5 @@
+import time
+
 import pypdfium2 # Needs to be at the top to avoid warnings
 import os
 os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # For some reason, transformers decided to use .isin for a simple op, which is not supported on MPS
@@ -20,18 +22,22 @@ def main():
     parser.add_argument("--start_page", type=int, default=None, help="Page to start processing at")
     parser.add_argument("--langs", type=str, help="Languages to use for OCR, comma separated", default=None)
     parser.add_argument("--batch_multiplier", type=int, default=2, help="How much to increase batch sizes")
+    parser.add_argument("--debug", action="store_true", help="Enable debug logging", default=False)
     args = parser.parse_args()
 
     langs = args.langs.split(",") if args.langs else None
 
     fname = args.filename
     model_lst = load_all_models()
+    start = time.time()
     full_text, images, out_meta = convert_single_pdf(fname, model_lst, max_pages=args.max_pages, langs=langs, batch_multiplier=args.batch_multiplier, start_page=args.start_page)
 
     fname = os.path.basename(fname)
     subfolder_path = save_markdown(args.output, fname, full_text, images, out_meta)
 
     print(f"Saved markdown to the {subfolder_path} folder")
+    if args.debug:
+        print(f"Total time: {time.time() - start}")
 
 
 if __name__ == "__main__":

marker/models.py

@@ -3,7 +3,7 @@ os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" # For some reason, transformers
 
 
 from marker.postprocessors.editor import load_editing_model
-from surya.model.detection import segformer
+from surya.model.detection.model import load_model as load_detection_model, load_processor as load_detection_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
@@ -25,11 +25,11 @@ def setup_recognition_model(langs, device=None, dtype=None):
 
 def setup_detection_model(device=None, dtype=None):
     if device:
-        model = segformer.load_model(device=device, dtype=dtype)
+        model = load_detection_model(device=device, dtype=dtype)
     else:
-        model = segformer.load_model()
+        model = load_detection_model()
 
-    processor = segformer.load_processor()
+    processor = load_detection_processor()
     model.processor = processor
     return model
 
@@ -46,10 +46,10 @@ def setup_texify_model(device=None, dtype=None):
 
 def setup_layout_model(device=None, dtype=None):
     if device:
-        model = segformer.load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT, device=device, dtype=dtype)
+        model = load_detection_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT, device=device, dtype=dtype)
     else:
-        model = segformer.load_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
-    processor = segformer.load_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+        model = load_detection_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
+    processor = load_detection_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
     model.processor = processor
     return model
 

marker/settings.py

@@ -67,7 +67,7 @@ class Settings(BaseSettings):
     # Layout model
     SURYA_LAYOUT_DPI: int = 96
     BAD_SPAN_TYPES: List[str] = ["Caption", "Footnote", "Page-footer", "Page-header", "Picture"]
-    LAYOUT_MODEL_CHECKPOINT: str = "vikp/surya_layout2"
+    LAYOUT_MODEL_CHECKPOINT: str = "vikp/surya_layout3"
     BBOX_INTERSECTION_THRESH: float = 0.7 # How much the layout and pdf bboxes need to overlap to be the same
     LAYOUT_BATCH_SIZE: Optional[int] = None # Defaults to 12 for cuda, 6 otherwise
 
@@ -83,10 +83,6 @@ class Settings(BaseSettings):
     ENABLE_EDITOR_MODEL: bool = False # The editor model can create false positives
     EDITOR_CUTOFF_THRESH: float = 0.9 # Ignore predictions below this probability
 
-    # Ray
-    RAY_CACHE_PATH: Optional[str] = None # Where to save ray cache
-    RAY_CORES_PER_WORKER: int = 1 # How many cpu cores to allocate per worker
-
     # Debug
     DEBUG: bool = False # Enable debug logging
     DEBUG_DATA_FOLDER: Optional[str] = None

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "marker-pdf"
-version = "0.2.15"
+version = "0.2.16"
 description = "Convert PDF to markdown with high speed and accuracy."
 authors = ["Vik Paruchuri <github@vikas.sh>"]
 readme = "README.md"
@@ -19,7 +19,7 @@ include = [
 
 [tool.poetry.dependencies]
 python = ">=3.9,<3.13,!=3.9.7"
-scikit-learn = "^1.3.2"
+scikit-learn = "^1.3.2,<=1.4.2"
 Pillow = "^10.1.0"
 pydantic = "^2.4.2"
 pydantic-settings = "^2.0.3"
@@ -32,7 +32,7 @@ tabulate = "^0.9.0"
 ftfy = "^6.1.1"
 texify = "^0.1.10"
 rapidfuzz = "^3.8.1"
-surya-ocr = "^0.4.14"
+surya-ocr = "^0.4.15"
 filetype = "^1.2.0"
 regex = "^2024.4.28"
 pdftext = "^0.3.10"