feat: new torch-based docling models (#120)

--------- Signed-off-by: Maxim Lysak <[email protected]> Co-authored-by: Maxim Lysak <[email protected]>
DS4SD · Oct 3, 2024 · 2422f70 · 2422f70
1 parent 9ebbbc1
commit 2422f70
Show file tree

Hide file tree

Showing 30 changed files with 1,157 additions and 1,183 deletions.
diff --git a/docling/document_converter.py b/docling/document_converter.py
@@ -67,7 +67,10 @@ def download_models_hf(
         from huggingface_hub import snapshot_download
 
         download_path = snapshot_download(
-            repo_id="ds4sd/docling-models", force_download=force, local_dir=local_dir
+            repo_id="ds4sd/docling-models",
+            force_download=force,
+            local_dir=local_dir,
+            revision="v2.0.0",
         )
 
         return Path(download_path)

diff --git a/docling/models/layout_model.py b/docling/models/layout_model.py
@@ -33,6 +33,7 @@ class LayoutModel:
         "Page-footer",
         "Code",
         "List-item",
+        # "Title"
         # "Formula",
     ]
     PAGE_HEADER_LABELS = ["Page-header", "Page-footer"]
@@ -69,9 +70,7 @@ def postprocess(self, clusters: List[Cluster], cells: List[Cell], page_height):
             "Key-Value Region": 0.45,
         }
 
-        CLASS_REMAPPINGS = {
-            "Document Index": "Table",
-        }
+        CLASS_REMAPPINGS = {"Document Index": "Table", "Title": "Section-header"}
 
         _log.debug("================= Start postprocess function ====================")
         start_time = time.time()
@@ -277,6 +276,7 @@ def __call__(self, page_batch: Iterable[Page]) -> Iterable[Page]:
                     bbox=BoundingBox.model_validate(pred_item),
                     cells=[],
                 )
+
                 clusters.append(cluster)
 
             # Map cells to clusters

diff --git a/docling/pipeline/standard_model_pipeline.py b/docling/pipeline/standard_model_pipeline.py
@@ -8,7 +8,7 @@
 
 
 class StandardModelPipeline(BaseModelPipeline):
-    _layout_model_path = "model_artifacts/layout/beehive_v0.0.5"
+    _layout_model_path = "model_artifacts/layout/beehive_v0.0.5_pt"
     _table_model_path = "model_artifacts/tableformer"
 
     def __init__(self, artifacts_path: Path, pipeline_options: PipelineOptions):

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,7 +38,7 @@ torchvision = [
 python = "^3.10"
 pydantic = "^2.0.0"
 docling-core = "^1.6.2"
-docling-ibm-models = "^1.3.1"
+docling-ibm-models = "^2.0.0"
 deepsearch-glm = "^0.22.0"
 filetype = "^1.2.0"
 pypdfium2 = "^4.30.0"