From 65bd8dec2ec2b52094e31dfb8609b8e2bae85772 Mon Sep 17 00:00:00 2001 From: "BBC, Esquire" Date: Wed, 14 Oct 2026 10:30:17 -0400 Subject: [PATCH] feat(asr): prefetch the native Distil-Whisper presets with docling-tools Add whisper_distil_small_en, whisper_distil_medium_en, whisper_distil_large_v3 and whisper_distil_large_v3_5 to docling-tools models download. They fetch the OpenAI-format checkpoints that the native Distil presets load, into the output directory's Hugging Face cache layout, which is the layout the native transcriber already reads when artifacts_path is set. Plain --all still excludes them, and --all --asr includes them. Move _DISTIL_WHISPER_OPENAI_CHECKPOINTS from the transcriber to asr_model_specs, next to the presets it describes, so the downloader and the transcriber share one table. docling.utils already depends on vlm_model_specs for the same reason, so tach.toml gets the matching asr_model_specs edges. The missing-checkpoint error under artifacts_path now gives the docling-tools command instead of hf download. Signed-off-by: BBC, Esquire --- .../.agents/skills/docling/references/cli.md | 5 ++- docling/cli/models.py | 32 +++++++++++++++++-- docling/datamodel/asr_model_specs.py | 12 +++++++ docling/pipeline/asr_transcriber.py | 20 ++++-------- docling/utils/model_downloader.py | 30 +++++++++++++++++ docs/usage/advanced_options.md | 12 ++++++- tach.toml | 2 ++ tests/test_asr_pipeline.py | 5 ++- tests/test_cli_tools.py | 22 ++++++++++++- 9 files changed, 121 insertions(+), 19 deletions(-) diff --git a/docling/.agents/skills/docling/references/cli.md b/docling/.agents/skills/docling/references/cli.md index dd46e8e8d7..c432f5cd3b 100644 --- a/docling/.agents/skills/docling/references/cli.md +++ b/docling/.agents/skills/docling/references/cli.md @@ -156,11 +156,14 @@ docling report.pdf --artifacts-path /models --output /tmp/ ``` For native Whisper ASR, prefetch a selected OpenAI-format checkpoint. Plain -`--all` excludes ASR models; `--all --asr` includes all native Whisper sizes. +`--all` excludes ASR models; `--all --asr` includes all native Whisper sizes and +the Distil-Whisper presets. ```bash docling-tools models download whisper_tiny --output-dir /models docling --asr-model whisper_tiny_native --artifacts-path /models audio.wav +docling-tools models download whisper_distil_small_en --output-dir /models +docling --asr-model whisper_distil_small_en_native --artifacts-path /models audio.wav ``` `--artifacts-path` overrides the default Hugging Face cache; you can also set diff --git a/docling/cli/models.py b/docling/cli/models.py index c88d5ce2d1..6c7da15a27 100644 --- a/docling/cli/models.py +++ b/docling/cli/models.py @@ -28,6 +28,10 @@ print(" pip install typer rich\n", file=sys.stderr) sys.exit(1) +from docling.datamodel.asr_model_specs import ( + _DISTIL_WHISPER_OPENAI_CHECKPOINTS, + _distil_whisper_prefetch_name, +) from docling.datamodel.settings import settings from docling.exceptions import RapidOcrModelSizeNotSupportedError from docling.models.stages.ocr.easyocr_model import ( @@ -85,6 +89,10 @@ class _AvailableModels(str, Enum): WHISPER_MEDIUM_EN = "whisper_medium_en" WHISPER_LARGE = "whisper_large" WHISPER_TURBO = "whisper_turbo" + WHISPER_DISTIL_SMALL_EN = "whisper_distil_small_en" + WHISPER_DISTIL_MEDIUM_EN = "whisper_distil_medium_en" + WHISPER_DISTIL_LARGE_V3 = "whisper_distil_large_v3" + WHISPER_DISTIL_LARGE_V3_5 = "whisper_distil_large_v3_5" _WHISPER_MODELS = { @@ -100,6 +108,11 @@ class _AvailableModels(str, Enum): _AvailableModels.WHISPER_TURBO: "turbo", } +_WHISPER_DISTIL_MODELS = { + _AvailableModels(_distil_whisper_prefetch_name(model_name)): model_name + for model_name in _DISTIL_WHISPER_OPENAI_CHECKPOINTS +} + _default_models = [ _AvailableModels.LAYOUT, @@ -142,7 +155,12 @@ def download( asr: Annotated[ bool, typer.Option( - ..., "--asr", help="With --all, also download native Whisper ASR models." + ..., + "--asr", + help=( + "With --all, also download native Whisper ASR models, including " + "the Distil-Whisper presets." + ), ), ] = False, quiet: Annotated[ @@ -213,7 +231,12 @@ def download( handlers=[RichHandler(show_level=False, show_time=False, markup=True)], ) to_download = models or ( - [model for model in _AvailableModels if asr or model not in _WHISPER_MODELS] + [ + model + for model in _AvailableModels + if asr + or (model not in _WHISPER_MODELS and model not in _WHISPER_DISTIL_MODELS) + ] if all else _default_models ) @@ -296,6 +319,11 @@ def download( whisper_models=[ _WHISPER_MODELS[model] for model in to_download if model in _WHISPER_MODELS ], + whisper_distil_models=[ + _WHISPER_DISTIL_MODELS[model] + for model in to_download + if model in _WHISPER_DISTIL_MODELS + ], ) if quiet: diff --git a/docling/datamodel/asr_model_specs.py b/docling/datamodel/asr_model_specs.py index 5133ef7aa0..6ea1f48b5a 100644 --- a/docling/datamodel/asr_model_specs.py +++ b/docling/datamodel/asr_model_specs.py @@ -533,6 +533,18 @@ def _get_whisper_turbo_model(): max_time_chunk=30.0, ) +_DISTIL_WHISPER_OPENAI_CHECKPOINTS: dict[str, tuple[str, str]] = { + "distil-small.en": ("distil-whisper/distil-small.en", "original-model.bin"), + "distil-medium.en": ("distil-whisper/distil-medium.en", "original-model.bin"), + "distil-large-v3": ("distil-whisper/distil-large-v3-openai", "model.bin"), + "distil-large-v3.5": ("distil-whisper/distil-large-v3.5-openai", "model.bin"), +} + + +def _distil_whisper_prefetch_name(model_name: str) -> str: + return "whisper_" + model_name.replace("-", "_").replace(".", "_") + + # ============================================================================= # WhisperS2T Models (CTranslate2 backend - fastest option for CPU/CUDA) # ============================================================================= diff --git a/docling/pipeline/asr_transcriber.py b/docling/pipeline/asr_transcriber.py index 07827807e0..e312f98e34 100644 --- a/docling/pipeline/asr_transcriber.py +++ b/docling/pipeline/asr_transcriber.py @@ -23,6 +23,10 @@ from docling.datamodel.accelerator_options import ( AcceleratorOptions, ) +from docling.datamodel.asr_model_specs import ( + _DISTIL_WHISPER_OPENAI_CHECKPOINTS, + _distil_whisper_prefetch_name, +) from docling.datamodel.base_models import ( ConversionStatus, DoclingComponentType, @@ -180,17 +184,6 @@ def to_string(self) -> str: return result -# Distil-Whisper models are not part of openai-whisper's model registry, but -# their Hugging Face repos publish the checkpoint in the original OpenAI -# format, which whisper.load_model() accepts as a local file path. -_DISTIL_WHISPER_OPENAI_CHECKPOINTS: dict[str, tuple[str, str]] = { - "distil-small.en": ("distil-whisper/distil-small.en", "original-model.bin"), - "distil-medium.en": ("distil-whisper/distil-medium.en", "original-model.bin"), - "distil-large-v3": ("distil-whisper/distil-large-v3-openai", "model.bin"), - "distil-large-v3.5": ("distil-whisper/distil-large-v3.5-openai", "model.bin"), -} - - class _NativeWhisperModel: def __init__( self, @@ -256,8 +249,9 @@ def __init__( f"artifacts_path ({artifacts_path}) does not contain " f"the checkpoint {repo_id}/{filename} required by ASR " f"model '{self.model_name}'. Prefetch it with: " - f"hf download {repo_id} {filename} " - f'--cache-dir "{artifacts_path}"' + f"docling-tools models download " + f"{_distil_whisper_prefetch_name(self.model_name)} " + f'--output-dir "{artifacts_path}"' ) from err else: checkpoint_path = hf_hub_download( diff --git a/docling/utils/model_downloader.py b/docling/utils/model_downloader.py index 523f5ed107..67fae1e243 100644 --- a/docling/utils/model_downloader.py +++ b/docling/utils/model_downloader.py @@ -10,6 +10,7 @@ import requests from tqdm import tqdm +from docling.datamodel.asr_model_specs import _DISTIL_WHISPER_OPENAI_CHECKPOINTS from docling.datamodel.pipeline_options import ( LayoutObjectDetectionOptions, granite_picture_description, @@ -121,6 +122,29 @@ def _sha256_file(path: Path) -> str: return digest.hexdigest() +def _download_distil_whisper_models( + model_names: list[str], output_dir: Path, *, force: bool, progress: bool +) -> None: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import disable_progress_bars + + unknown = set(model_names) - _DISTIL_WHISPER_OPENAI_CHECKPOINTS.keys() + if unknown: + raise ValueError(f"Unknown Distil-Whisper models: {sorted(unknown)}") + if not progress: + disable_progress_bars() + + for model_name in model_names: + repo_id, filename = _DISTIL_WHISPER_OPENAI_CHECKPOINTS[model_name] + _log.info("Downloading Distil-Whisper model %s...", model_name) + hf_hub_download( + repo_id=repo_id, + filename=filename, + cache_dir=output_dir, + force_download=force, + ) + + def download_models( output_dir: Optional[Path] = None, *, @@ -148,6 +172,7 @@ def download_models( easyocr_languages: Optional[list[str]] = None, # BCP-47 tags with_nemotron_ocr: bool = False, whisper_models: Optional[list[str]] = None, + whisper_distil_models: Optional[list[str]] = None, ): if easyocr_languages is not None and not with_easyocr: raise ValueError("easyocr_languages requires with_easyocr=True") @@ -171,6 +196,11 @@ def download_models( whisper_models, output_dir, force=force, progress=progress ) + if whisper_distil_models: + _download_distil_whisper_models( + whisper_distil_models, output_dir, force=force, progress=progress + ) + if with_layout: _log.info("Downloading layout model...") layout_spec = LayoutObjectDetectionOptions().model_spec diff --git a/docs/usage/advanced_options.md b/docs/usage/advanced_options.md index cee4f3eca5..b144660249 100644 --- a/docs/usage/advanced_options.md +++ b/docs/usage/advanced_options.md @@ -36,8 +36,18 @@ docling-tools models download whisper_tiny --output-dir /models docling --asr-model whisper_tiny_native --artifacts-path /models audio.wav ``` +The native Distil-Whisper presets work the same way. Their OpenAI-format +checkpoints come from Hugging Face and are stored in the output directory in the +Hugging Face cache layout: + +```sh +docling-tools models download whisper_distil_small_en --output-dir /models +docling --asr-model whisper_distil_small_en_native --artifacts-path /models audio.wav +``` + `docling-tools models download --all` does not include ASR models. Add `--asr` -to prefetch all native Whisper model sizes and English-only variants: +to prefetch all native Whisper model sizes, English-only variants, and +Distil-Whisper presets: ```sh docling-tools models download --all --asr --output-dir /models diff --git a/tach.toml b/tach.toml index 0b2160d683..ac3fb75e02 100644 --- a/tach.toml +++ b/tach.toml @@ -186,6 +186,7 @@ path = "docling.utils" layer = "core" depends_on = [ "docling.datamodel", + "docling.datamodel.asr_model_specs", "docling.datamodel.vlm_model_specs", "docling.models.stages.chart_extraction", "docling.models.stages.code_formula", @@ -463,6 +464,7 @@ layer = "pipeline" depends_on = [ "docling.backend", "docling.datamodel", + "docling.datamodel.asr_model_specs", "docling.datamodel.pipeline_options_asr_model", "docling.pipeline.base_pipeline", "docling.utils", diff --git a/tests/test_asr_pipeline.py b/tests/test_asr_pipeline.py index 374334dde1..6114d2a5ce 100644 --- a/tests/test_asr_pipeline.py +++ b/tests/test_asr_pipeline.py @@ -278,7 +278,10 @@ def test_native_distil_artifacts_path_missing_checkpoint_raises(tmp_path): max_time_chunk=1.0, language="en", ) - with pytest.raises(FileNotFoundError, match="does not contain"): + with pytest.raises( + FileNotFoundError, + match=r"does not contain .* docling-tools models download whisper_distil_small_en", + ): _NativeWhisperModel( True, tmp_path, AcceleratorOptions(device=AcceleratorDevice.CPU), opts ) diff --git a/tests/test_cli_tools.py b/tests/test_cli_tools.py index 62f4a8b5b8..b7d6014cbb 100644 --- a/tests/test_cli_tools.py +++ b/tests/test_cli_tools.py @@ -20,6 +20,7 @@ import docling.cli.models as models_cli from docling.cli.models import _AvailableModels, _default_models from docling.cli.tools import app +from docling.datamodel.asr_model_specs import _DISTIL_WHISPER_OPENAI_CHECKPOINTS # Under CI Rich thinks it has a terminal and styles the error panel, landing # escapes between the border and the wrapped halves of a sentence. @@ -121,11 +122,16 @@ def test_download_all_selects_every_available_model(tmp_path, recorded_download) assert result.exit_code == 0 selection = {k: v for k, v in recorded_download.items() if k.startswith("with_")} - assert len(selection) == len(_AvailableModels) - len(models_cli._WHISPER_MODELS) + assert len(selection) == ( + len(_AvailableModels) + - len(models_cli._WHISPER_MODELS) + - len(models_cli._WHISPER_DISTIL_MODELS) + ) assert all(selection.values()), ( f"not selected by --all: {_enabled(recorded_download) ^ set(selection)}" ) assert recorded_download["whisper_models"] == [] + assert recorded_download["whisper_distil_models"] == [] def test_download_all_with_asr_includes_native_whisper(tmp_path, recorded_download): @@ -137,6 +143,9 @@ def test_download_all_with_asr_includes_native_whisper(tmp_path, recorded_downlo assert set(recorded_download["whisper_models"]) == set( models_cli._WHISPER_MODELS.values() ) + assert set(recorded_download["whisper_distil_models"]) == set( + _DISTIL_WHISPER_OPENAI_CHECKPOINTS + ) def test_download_one_native_whisper_model(tmp_path, recorded_download): @@ -149,6 +158,17 @@ def test_download_one_native_whisper_model(tmp_path, recorded_download): assert _enabled(recorded_download) == set() +def test_download_one_distil_whisper_model(tmp_path, recorded_download): + result = runner.invoke( + app, ["models", "download", "-o", str(tmp_path), "whisper_distil_small_en"] + ) + + assert result.exit_code == 0 + assert recorded_download["whisper_distil_models"] == ["distil-small.en"] + assert recorded_download["whisper_models"] == [] + assert _enabled(recorded_download) == set() + + def test_download_asr_requires_all(tmp_path): result = runner.invoke(app, ["models", "download", "-o", str(tmp_path), "--asr"])