diff --git a/docling/.agents/skills/docling/references/cli.md b/docling/.agents/skills/docling/references/cli.md index 5408deab20..fc3aa691fc 100644 --- a/docling/.agents/skills/docling/references/cli.md +++ b/docling/.agents/skills/docling/references/cli.md @@ -178,13 +178,16 @@ 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 and -the Distil-Whisper presets. +the Distil-Whisper presets. WhisperS2T presets are prefetched only by name; each +one downloads the `float16` and `float32` models and the word aligner it needs. ```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 +docling-tools models download whisper_small_en_s2t --output-dir /models +docling --asr-model whisper_small_en_s2t --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 6c7da15a27..3475c629c8 100644 --- a/docling/cli/models.py +++ b/docling/cli/models.py @@ -30,7 +30,9 @@ from docling.datamodel.asr_model_specs import ( _DISTIL_WHISPER_OPENAI_CHECKPOINTS, + _WHISPER_S2T_PRESETS, _distil_whisper_prefetch_name, + _whisper_s2t_prefetch_name, ) from docling.datamodel.settings import settings from docling.exceptions import RapidOcrModelSizeNotSupportedError @@ -93,6 +95,20 @@ class _AvailableModels(str, Enum): 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_TINY_S2T = "whisper_tiny_s2t" + WHISPER_TINY_EN_S2T = "whisper_tiny_en_s2t" + WHISPER_BASE_S2T = "whisper_base_s2t" + WHISPER_BASE_EN_S2T = "whisper_base_en_s2t" + WHISPER_SMALL_S2T = "whisper_small_s2t" + WHISPER_SMALL_EN_S2T = "whisper_small_en_s2t" + WHISPER_DISTIL_SMALL_EN_S2T = "whisper_distil_small_en_s2t" + WHISPER_MEDIUM_S2T = "whisper_medium_s2t" + WHISPER_MEDIUM_EN_S2T = "whisper_medium_en_s2t" + WHISPER_DISTIL_MEDIUM_EN_S2T = "whisper_distil_medium_en_s2t" + WHISPER_LARGE_V3_S2T = "whisper_large_v3_s2t" + WHISPER_DISTIL_LARGE_V3_S2T = "whisper_distil_large_v3_s2t" + WHISPER_DISTIL_LARGE_V3_5_S2T = "whisper_distil_large_v3_5_s2t" + WHISPER_LARGE_V3_TURBO_S2T = "whisper_large_v3_turbo_s2t" _WHISPER_MODELS = { @@ -113,6 +129,11 @@ class _AvailableModels(str, Enum): for model_name in _DISTIL_WHISPER_OPENAI_CHECKPOINTS } +_WHISPER_S2T_MODELS = { + _AvailableModels(_whisper_s2t_prefetch_name(model_name)): model_name + for model_name in _WHISPER_S2T_PRESETS +} + _default_models = [ _AvailableModels.LAYOUT, @@ -159,7 +180,8 @@ def download( "--asr", help=( "With --all, also download native Whisper ASR models, including " - "the Distil-Whisper presets." + "the Distil-Whisper presets. WhisperS2T models are downloaded " + "only by name." ), ), ] = False, @@ -234,8 +256,13 @@ def download( [ model for model in _AvailableModels - if asr - or (model not in _WHISPER_MODELS and model not in _WHISPER_DISTIL_MODELS) + if model not in _WHISPER_S2T_MODELS + and ( + asr + or ( + model not in _WHISPER_MODELS and model not in _WHISPER_DISTIL_MODELS + ) + ) ] if all else _default_models @@ -324,6 +351,11 @@ def download( for model in to_download if model in _WHISPER_DISTIL_MODELS ], + whisper_s2t_models=[ + _WHISPER_S2T_MODELS[model] + for model in to_download + if model in _WHISPER_S2T_MODELS + ], ) if quiet: diff --git a/docling/datamodel/asr_model_specs.py b/docling/datamodel/asr_model_specs.py index 4609754881..635c2a85b7 100644 --- a/docling/datamodel/asr_model_specs.py +++ b/docling/datamodel/asr_model_specs.py @@ -718,6 +718,43 @@ def _distil_whisper_prefetch_name(model_name: str) -> str: beam_size=1, ) +_WHISPER_S2T_PRESETS: dict[str, InlineAsrWhisperS2TOptions] = { + preset.repo_id: preset + for preset in ( + WHISPER_TINY_S2T, + WHISPER_TINY_EN_S2T, + WHISPER_BASE_S2T, + WHISPER_BASE_EN_S2T, + WHISPER_SMALL_S2T, + WHISPER_SMALL_EN_S2T, + WHISPER_DISTIL_SMALL_EN_S2T, + WHISPER_MEDIUM_S2T, + WHISPER_MEDIUM_EN_S2T, + WHISPER_DISTIL_MEDIUM_EN_S2T, + WHISPER_LARGE_V3_S2T, + WHISPER_DISTIL_LARGE_V3_S2T, + WHISPER_DISTIL_LARGE_V3_5_S2T, + WHISPER_LARGE_V3_TURBO_S2T, + ) +} + +_WHISPER_S2T_COMPUTE_TYPES = ("float32", "float16", "bfloat16") + + +def _whisper_s2t_compute_type(torch_dtype: str | None, device: str) -> str: + if ( + device == "cpu" + or torch_dtype is None + or torch_dtype not in _WHISPER_S2T_COMPUTE_TYPES + ): + return "float32" + return torch_dtype + + +def _whisper_s2t_prefetch_name(model_name: str) -> str: + return "whisper_" + model_name.replace("-", "_").replace(".", "_") + "_s2t" + + # ============================================================================= # Note on auto-selecting models # ============================================================================= diff --git a/docling/pipeline/asr_transcriber.py b/docling/pipeline/asr_transcriber.py index 8203c29713..8f11372449 100644 --- a/docling/pipeline/asr_transcriber.py +++ b/docling/pipeline/asr_transcriber.py @@ -28,7 +28,10 @@ ) from docling.datamodel.asr_model_specs import ( _DISTIL_WHISPER_OPENAI_CHECKPOINTS, + _WHISPER_S2T_PRESETS, _distil_whisper_prefetch_name, + _whisper_s2t_compute_type, + _whisper_s2t_prefetch_name, ) from docling.datamodel.base_models import ( ConversionStatus, @@ -669,7 +672,7 @@ def __init__( except ImportError: raise ImportError( "whisper_s2t is not installed. Please install it via " - "`pip install 'whisper-s2t-reborn[pyav]>=1.7.1'`." + "`pip install 'whisper-s2t-reborn>=1.8.0'`." ) self.whisper_s2t = whisper_s2t @@ -692,13 +695,15 @@ def __init__( # inference. Coerce to float32 when running on CPU so that the # explicit *_S2T presets (which default to float16 for CUDA # performance) do not fail at model load on CPU-only installs. - compute_type = asr_options.torch_dtype - if self.device == "cpu" and compute_type in ("float16", "bfloat16"): + compute_type = _whisper_s2t_compute_type( + asr_options.torch_dtype, self.device + ) + if asr_options.torch_dtype not in (None, compute_type): _log.warning( - f"compute_type='{compute_type}' is not supported by " - f"CTranslate2 on CPU; falling back to 'float32'." + f"compute_type='{asr_options.torch_dtype}' is not supported " + f"by WhisperS2T on {self.device}; falling back to " + f"'{compute_type}'." ) - compute_type = "float32" # Build ASR options for whisper_s2t asr_opts = { @@ -715,10 +720,34 @@ def __init__( "asr_options": asr_opts, } - self.model = whisper_s2t.load_model( - model_identifier=self.model_identifier, - **model_kwargs, - ) + if artifacts_path is None: + self.model = whisper_s2t.load_model( + model_identifier=self.model_identifier, + **model_kwargs, + ) + else: + from huggingface_hub.utils import LocalEntryNotFoundError + + try: + self.model = whisper_s2t.load_model( + model_identifier=self.model_identifier, + cache_dir=str(artifacts_path), + local_files_only=True, + **model_kwargs, + ) + except LocalEntryNotFoundError as err: + message = ( + f"artifacts_path ({artifacts_path}) does not contain the " + f"{compute_type} WhisperS2T files required by ASR model " + f"'{self.model_identifier}'." + ) + if self.model_identifier in _WHISPER_S2T_PRESETS: + message += ( + " Prefetch them with: docling-tools models download " + f"{_whisper_s2t_prefetch_name(self.model_identifier)} " + f'--output-dir "{artifacts_path}"' + ) + raise FileNotFoundError(message) from err # Store options for transcription self.language = asr_options.language diff --git a/docling/utils/model_downloader.py b/docling/utils/model_downloader.py index dab692ad68..8880686a77 100644 --- a/docling/utils/model_downloader.py +++ b/docling/utils/model_downloader.py @@ -10,7 +10,11 @@ import requests from tqdm import tqdm -from docling.datamodel.asr_model_specs import _DISTIL_WHISPER_OPENAI_CHECKPOINTS +from docling.datamodel.asr_model_specs import ( + _DISTIL_WHISPER_OPENAI_CHECKPOINTS, + _WHISPER_S2T_PRESETS, + _whisper_s2t_compute_type, +) from docling.datamodel.pipeline_options import ( LayoutObjectDetectionOptions, granite_picture_description, @@ -146,6 +150,47 @@ def _download_distil_whisper_models( ) +def _download_whisper_s2t_models( + model_names: list[str], output_dir: Path, *, force: bool, progress: bool +) -> None: + try: + import whisper_s2t # type: ignore + except ImportError as exc: + raise ImportError( + "WhisperS2T prefetch requires whisper-s2t-reborn. Install it via " + "`pip install 'whisper-s2t-reborn>=1.8.0'`." + ) from exc + from huggingface_hub.utils import disable_progress_bars + + unknown = set(model_names) - _WHISPER_S2T_PRESETS.keys() + if unknown: + raise ValueError(f"Unknown WhisperS2T models: {sorted(unknown)}") + if not progress: + disable_progress_bars() + + for model_name in model_names: + torch_dtype = _WHISPER_S2T_PRESETS[model_name].torch_dtype + compute_types = sorted( + { + _whisper_s2t_compute_type(torch_dtype, device) + for device in ("cuda", "cpu") + } + ) + aligner = whisper_s2t.default_word_aligner(model_name) + names = [model_name] if aligner is None else [model_name, aligner] + for compute_type in compute_types: + _log.info( + "Downloading WhisperS2T model %s (%s)...", model_name, compute_type + ) + for name in names: + whisper_s2t.download_model( + name, + compute_type=compute_type, + cache_dir=str(output_dir), + force_download=force, + ) + + def download_models( output_dir: Optional[Path] = None, *, @@ -174,6 +219,7 @@ def download_models( with_nemotron_ocr: bool = False, whisper_models: Optional[list[str]] = None, whisper_distil_models: Optional[list[str]] = None, + whisper_s2t_models: Optional[list[str]] = None, ): if easyocr_languages is not None and not with_easyocr: raise ValueError("easyocr_languages requires with_easyocr=True") @@ -202,6 +248,11 @@ def download_models( whisper_distil_models, output_dir, force=force, progress=progress ) + if whisper_s2t_models: + _download_whisper_s2t_models( + whisper_s2t_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 b79bc1dc3d..7cb1356154 100644 --- a/docs/usage/advanced_options.md +++ b/docs/usage/advanced_options.md @@ -45,6 +45,16 @@ docling-tools models download whisper_distil_small_en --output-dir /models docling --asr-model whisper_distil_small_en_native --artifacts-path /models audio.wav ``` +The WhisperS2T presets use their `--asr-model` names. Each one downloads its +CTranslate2 model in `float16` (for CUDA) and `float32` (for CPU), plus the +small word-alignment model that word timestamps use if the preset needs one, +into the output directory in the Hugging Face cache layout: + +```sh +docling-tools models download whisper_small_en_s2t --output-dir /models +docling --asr-model whisper_small_en_s2t --artifacts-path /models audio.wav +``` + `docling-tools models download --all` does not include ASR models. Add `--asr` to prefetch all native Whisper model sizes, English-only variants, and Distil-Whisper presets: @@ -54,7 +64,8 @@ docling-tools models download --all --asr --output-dir /models ``` This requires the `asr` extra and downloads OpenAI-format checkpoints. It does -not prefetch Transformers-format, MLX, or WhisperS2T models. +not prefetch Transformers-format, MLX, or WhisperS2T models. Prefetch a +WhisperS2T preset by its name, as shown above. Alternatively, models can be programmatically downloaded using `docling.utils.model_downloader.download_models()`. diff --git a/pyproject.toml b/pyproject.toml index 12bb71482c..8a73a417fe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -181,7 +181,7 @@ format-audio = [ 'mlx-whisper>=0.4.3 ; python_version >= "3.10" and sys_platform == "darwin" and platform_machine == "arm64"', 'openai-whisper>=20250625', 'numba>=0.63.0', - 'whisper-s2t-reborn>=1.7.1,<2.0.0 ; sys_platform != "darwin" or platform_machine != "arm64"', + 'whisper-s2t-reborn>=1.8.0,<2.0.0 ; sys_platform != "darwin" or platform_machine != "arm64"', ] feat-asr-transformers = [ diff --git a/tests/test_asr_whisper_s2t.py b/tests/test_asr_whisper_s2t.py index 050e175c7f..b40a79e458 100644 --- a/tests/test_asr_whisper_s2t.py +++ b/tests/test_asr_whisper_s2t.py @@ -18,6 +18,7 @@ InlineAsrWhisperS2TOptions, ) from docling.pipeline.asr_pipeline import AsrPipeline, _WhisperS2TModel +from docling.utils.model_downloader import download_models pytestmark = pytest.mark.ml_asr @@ -560,9 +561,7 @@ def test_whisper_s2t_cpu_coerces_bfloat16_compute_type(self): f"got {call_kwargs.get('compute_type')!r}" ) - def test_whisper_s2t_cpu_preserves_cpu_compatible_compute_types(self): - """Regression test: compute_type values already supported by CTranslate2 - on CPU (float32, int8, int8_float32) must be passed through unchanged.""" + def test_whisper_s2t_cpu_runs_other_compute_types_as_float32(self): mock_whisper_s2t = Mock() mock_whisper_s2t.load_model.return_value = Mock() @@ -585,9 +584,9 @@ def test_whisper_s2t_cpu_preserves_cpu_compatible_compute_types(self): ) call_kwargs = mock_whisper_s2t.load_model.call_args.kwargs - assert call_kwargs.get("compute_type") == compute_type, ( - f"compute_type={compute_type!r} on CPU should pass through " - f"unchanged, got {call_kwargs.get('compute_type')!r}" + assert call_kwargs.get("compute_type") == "float32", ( + f"compute_type={compute_type!r} on CPU should run as " + f"'float32', got {call_kwargs.get('compute_type')!r}" ) mock_whisper_s2t.load_model.reset_mock() @@ -664,3 +663,57 @@ def test_asr_pipeline_with_whisper_s2t(self): pipeline = AsrPipeline(pipeline_options) assert isinstance(pipeline._model, _WhisperS2TModel) assert pipeline._model.model_identifier == "tiny" + + +class TestWhisperS2TOffline: + def test_artifacts_path_without_the_cpu_files_raises(self, tmp_path): + pytest.importorskip("whisper_s2t") + asr_options = InlineAsrWhisperS2TOptions( + repo_id="small.en", + inference_framework=InferenceAsrFramework.WHISPER_S2T, + language="en", + task="transcribe", + torch_dtype="float16", + ) + + with pytest.raises( + FileNotFoundError, + match=( + r"does not contain the float32 .* " + r"docling-tools models download whisper_small_en_s2t" + ), + ): + _WhisperS2TModel( + enabled=True, + artifacts_path=tmp_path, + accelerator_options=AcceleratorOptions(device=AcceleratorDevice.CPU), + asr_options=asr_options, + ) + + def test_prefetch_fetches_the_model_and_its_aligner_at_both_precisions( + self, tmp_path, monkeypatch + ): + whisper_s2t = pytest.importorskip("whisper_s2t") + fetched = [] + + def _download_model(name, compute_type, cache_dir, force_download): + fetched.append((name, compute_type, cache_dir, force_download)) + return str(tmp_path / name) + + monkeypatch.setattr(whisper_s2t, "download_model", _download_model) + download_models( + output_dir=tmp_path, + force=True, + with_layout=False, + with_tableformer=False, + with_code_formula=False, + with_picture_classifier=False, + with_rapidocr=False, + whisper_s2t_models=["small.en", "large-v3"], + ) + + assert sorted(fetched) == sorted( + (name, compute_type, str(tmp_path), True) + for name in ("small.en", "tiny.en", "large-v3") + for compute_type in ("float16", "float32") + ) diff --git a/tests/test_cli_tools.py b/tests/test_cli_tools.py index b7d6014cbb..c403cc6f71 100644 --- a/tests/test_cli_tools.py +++ b/tests/test_cli_tools.py @@ -126,12 +126,14 @@ def test_download_all_selects_every_available_model(tmp_path, recorded_download) len(_AvailableModels) - len(models_cli._WHISPER_MODELS) - len(models_cli._WHISPER_DISTIL_MODELS) + - len(models_cli._WHISPER_S2T_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"] == [] + assert recorded_download["whisper_s2t_models"] == [] def test_download_all_with_asr_includes_native_whisper(tmp_path, recorded_download): @@ -146,6 +148,7 @@ def test_download_all_with_asr_includes_native_whisper(tmp_path, recorded_downlo assert set(recorded_download["whisper_distil_models"]) == set( _DISTIL_WHISPER_OPENAI_CHECKPOINTS ) + assert recorded_download["whisper_s2t_models"] == [] def test_download_one_native_whisper_model(tmp_path, recorded_download): @@ -169,6 +172,18 @@ def test_download_one_distil_whisper_model(tmp_path, recorded_download): assert _enabled(recorded_download) == set() +def test_download_one_whisper_s2t_model(tmp_path, recorded_download): + result = runner.invoke( + app, ["models", "download", "-o", str(tmp_path), "whisper_small_en_s2t"] + ) + + assert result.exit_code == 0 + assert recorded_download["whisper_s2t_models"] == ["small.en"] + assert recorded_download["whisper_models"] == [] + assert recorded_download["whisper_distil_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"]) diff --git a/uv.lock b/uv.lock index 9df470be4d..dc898a3e82 100644 --- a/uv.lock +++ b/uv.lock @@ -2624,7 +2624,7 @@ requires-dist = [ { name = "typer", marker = "extra == 'cli'", specifier = ">=0.19.0,<0.27.0" }, { name = "typer", marker = "extra == 'service-client'", specifier = ">=0.19.0,<0.27.0" }, { name = "websockets", marker = "extra == 'service-client'", specifier = ">=14.0,<17.0" }, - { name = "whisper-s2t-reborn", marker = "(platform_machine != 'arm64' and extra == 'format-audio') or (sys_platform != 'darwin' and extra == 'format-audio')", specifier = ">=1.7.1,<2.0.0" }, + { name = "whisper-s2t-reborn", marker = "(platform_machine != 'arm64' and extra == 'format-audio') or (sys_platform != 'darwin' and extra == 'format-audio')", specifier = ">=1.8.0,<2.0.0" }, ] provides-extras = ["convert-core", "extract-core", "format-pdf-pypdfium2", "format-pdf-docling", "format-pdf", "format-docx", "format-pptx", "format-xlsx", "format-office", "format-opendocument", "format-html", "format-markdown", "format-web", "format-latex", "format-email", "format-iwork", "format-xml-jats", "format-xml-uspto", "format-xml-xbrl", "format-html-render", "format-audio", "feat-asr-transformers", "format-video", "feat-ocr-rapidocr", "feat-ocr-rapidocr-onnx", "feat-ocr-easyocr", "feat-ocr-tesserocr", "feat-ocr-mac", "feat-ocr-nemotron", "models-local", "models-remote", "models-onnxruntime", "models-vlm-inline", "feat-chunking", "service-client", "cli", "standard", "all"]