| import logging |
| from pathlib import Path |
| from typing import Optional |
|
|
| from docling.datamodel.pipeline_options import smolvlm_picture_description |
| 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.picture_description_vlm_model import PictureDescriptionVlmModel |
| 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_smolvlm: bool = True, |
| with_easyocr: bool = True, |
| ): |
| if output_dir is None: |
| output_dir = settings.cache_dir / "models" |
|
|
| |
| 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_smolvlm: |
| _log.info(f"Downloading SmolVlm model...") |
| PictureDescriptionVlmModel.download_models( |
| repo_id=smolvlm_picture_description.repo_id, |
| local_dir=output_dir / smolvlm_picture_description.repo_cache_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 |
|
|
