feat: new artifacts path and CLI utility (#876)

* fix artifacts path Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * add docling-models utility Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * missing formatting Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * rename utility to docling-tools Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * rename download methods and deprecation warnings Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * propagate artifacts path usage for ocr models Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * move function to utils Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * remove unused file Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * update docs Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> * simplify downloading specific model(s) Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com> * minor refactor Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com> --------- Signed-off-by: Michele Dolfi <dol@zurich.ibm.com> Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com> Co-authored-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
2025-12-12 14:48:21 +00:00 · 2025-02-06 15:46:32 +01:00
parent 722a6eb7b9
commit ed74fe2ec0
15 changed files with 467 additions and 68 deletions
--- a/docling/utils/model_downloader.py
+++ b/docling/utils/model_downloader.py
@@ -0,0 +1,72 @@
+import logging
+from pathlib import Path
+from typing import Optional
+
+from docling.datamodel.settings import settings
+from docling.models.code_formula_model import CodeFormulaModel
+from docling.models.document_picture_classifier import DocumentPictureClassifier
+from docling.models.easyocr_model import EasyOcrModel
+from docling.models.layout_model import LayoutModel
+from docling.models.table_structure_model import TableStructureModel
+
+_log = logging.getLogger(__name__)
+
+
+def download_models(
+    output_dir: Optional[Path] = None,
+    *,
+    force: bool = False,
+    progress: bool = False,
+    with_layout: bool = True,
+    with_tableformer: bool = True,
+    with_code_formula: bool = True,
+    with_picture_classifier: bool = True,
+    with_easyocr: bool = True,
+):
+    if output_dir is None:
+        output_dir = settings.cache_dir / "models"
+
+    # Make sure the folder exists
+    output_dir.mkdir(exist_ok=True, parents=True)
+
+    if with_layout:
+        _log.info(f"Downloading layout model...")
+        LayoutModel.download_models(
+            local_dir=output_dir / LayoutModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_tableformer:
+        _log.info(f"Downloading tableformer model...")
+        TableStructureModel.download_models(
+            local_dir=output_dir / TableStructureModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_picture_classifier:
+        _log.info(f"Downloading picture classifier model...")
+        DocumentPictureClassifier.download_models(
+            local_dir=output_dir / DocumentPictureClassifier._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_code_formula:
+        _log.info(f"Downloading code formula model...")
+        CodeFormulaModel.download_models(
+            local_dir=output_dir / CodeFormulaModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    if with_easyocr:
+        _log.info(f"Downloading easyocr models...")
+        EasyOcrModel.download_models(
+            local_dir=output_dir / EasyOcrModel._model_repo_folder,
+            force=force,
+            progress=progress,
+        )
+
+    return output_dir