# SPDX-FileCopyrightText: The Docling Contributors # SPDX-License-Identifier: MIT import logging import shutil import sys import tempfile from io import BytesIO from pathlib import Path from typing import Final, Protocol from docling_core.types.doc import ( ContentLayer, DocItemLabel, DoclingDocument, DocumentOrigin, TrackSource, ) from pydantic import BaseModel, Field from docling.backend.abstract_backend import AbstractDocumentBackend from docling.backend.noop_backend import NoOpBackend from docling.datamodel.accelerator_options import ( AcceleratorOptions, ) from docling.datamodel.base_models import ( ConversionStatus, DoclingComponentType, ErrorItem, ) from docling.datamodel.document import ConversionResult from docling.datamodel.pipeline_options import ( AsrPipelineOptions, ) from docling.datamodel.pipeline_options_asr_model import ( InlineAsrMlxWhisperOptions, InlineAsrNativeWhisperOptions, InlineAsrOptions, InlineAsrWhisperS2TOptions, ) from docling.pipeline.base_pipeline import BasePipeline from docling.utils.accelerator_utils import decide_device from docling.utils.profiling import ProfilingScope, TimeRecorder _log = logging.getLogger(__name__) ZERO_DURATION_SEGMENT_EPS: Final[float] = 0.001 """Minimal duration (in seconds) to add to zero-duration ASR segments. When an ASR segment has end_time <= start_time but contains non-empty text, this epsilon value is added to the start_time to create a valid time range. This prevents validation issues with Docling data models. """ MISSING_FFMPEG_MESSAGE: Final[str] = ( "FFmpeg is required for audio processing but was not found on PATH. " "Install it with your system package manager (e.g., 'brew install ffmpeg' " "on macOS, 'apt-get install ffmpeg' on Linux, 'winget install ffmpeg' on " "Windows)." ) _AUDIO_SUFFIX_TO_MIMETYPE = { ".wav": "audio/x-wav", ".mp3": "audio/mp3", ".m4a": "audio/m4a", ".aac": "audio/aac", ".ogg": "audio/ogg", ".flac": "audio/flac", } def _audio_mimetype(filename: str) -> str: suffix = Path(filename).suffix.lower() return _AUDIO_SUFFIX_TO_MIMETYPE.get(suffix, "audio/x-wav") def _process_conversation( conversation: list["_ConversationItem"], conv_res: ConversionResult ) -> None: """Process the conversation items and add them to the document.""" # Ensure we have a proper DoclingDocument filename = conv_res.input.file.name or "audio.wav" origin = DocumentOrigin( filename=filename, mimetype=_audio_mimetype(filename), binary_hash=conv_res.input.document_hash, ) conv_res.document = DoclingDocument( name=conv_res.input.file.stem or "audio.wav", origin=origin ) for citem in conversation: # Fix zero-duration segments (end_time <= start_time) with non-empty text if ( citem.start_time is not None and citem.end_time is not None and citem.end_time <= citem.start_time and citem.text.strip() ): _log.warning( f"Zero-duration ASR segment at {citem.start_time}s: " f"'{citem.text}' - adjusting end_time" ) citem.end_time = citem.start_time + ZERO_DURATION_SEGMENT_EPS # Add all segments with valid timestamps and non-empty text if ( citem.start_time is not None and citem.end_time is not None and citem.text.strip() ): try: track: TrackSource = TrackSource( start_time=citem.start_time, end_time=citem.end_time, voice=citem.speaker, ) _ = conv_res.document.add_text( label=DocItemLabel.TEXT, text=citem.text, content_layer=ContentLayer.BODY, source=track, ) except Exception as e: _log.warning( f"Failed to add conversation item to document " f"(start: {citem.start_time}s, end: {citem.end_time}s, " f"speaker: {citem.speaker}, text: '{citem.text[:50]}...'): " f"{e}. Skipping this item and continuing with the rest." ) continue class _ConversationWord(BaseModel): text: str start_time: float | None = Field( None, description="Start time in seconds from video start" ) end_time: float | None = Field( None, ge=0, description="End time in seconds from video start" ) class _ConversationItem(BaseModel): text: str start_time: float | None = Field( None, description="Start time in seconds from video start" ) end_time: float | None = Field( None, ge=0, description="End time in seconds from video start" ) speaker_id: int | None = Field(None, description="Numeric speaker identifier") speaker: str | None = Field( None, description="Speaker name, defaults to speaker-{speaker_id}" ) words: list[_ConversationWord] | None = Field( None, description="Individual words with time-stamps" ) def __lt__(self, other): if not isinstance(other, _ConversationItem): return NotImplemented return self.start_time < other.start_time def __eq__(self, other): if not isinstance(other, _ConversationItem): return NotImplemented return self.start_time == other.start_time def to_string(self) -> str: """Format the conversation entry as a string""" result = "" if (self.start_time is not None) and (self.end_time is not None): result += f"[time: {self.start_time}-{self.end_time}] " if self.speaker is not None: result += f"[speaker:{self.speaker}] " result += self.text 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, enabled: bool, artifacts_path: Path | None, accelerator_options: AcceleratorOptions, asr_options: InlineAsrNativeWhisperOptions, ): """Transcriber using native Whisper.""" self.enabled = enabled _log.info(f"artifacts-path: {artifacts_path}") _log.info(f"accelerator_options: {accelerator_options}") if self.enabled: try: import whisper # type: ignore except ImportError: if sys.version_info < (3, 14): raise ImportError( "whisper is not installed. Please install it via " "`pip install openai-whisper` or do `uv sync --extra asr`." ) else: raise ImportError( "whisper is not installed. Unfortunately its dependencies " "are not yet available for Python 3.14." ) self.asr_options = asr_options self.max_tokens = asr_options.max_new_tokens self.device = decide_device( accelerator_options.device, supported_devices=asr_options.supported_devices, ) _log.info(f"Available device for Whisper: {self.device}") self.model_name = asr_options.repo_id _log.info(f"loading _NativeWhisperModel({self.model_name})") distil_checkpoint = _DISTIL_WHISPER_OPENAI_CHECKPOINTS.get(self.model_name) if distil_checkpoint is not None: from huggingface_hub import hf_hub_download from huggingface_hub.utils import LocalEntryNotFoundError repo_id, filename = distil_checkpoint _log.info( f"loading {self.model_name} from OpenAI-format checkpoint " f"{repo_id}/{filename}" ) if artifacts_path is not None: # artifacts_path means fully-offline operation: resolve the # checkpoint from the local cache and never download. try: checkpoint_path = hf_hub_download( repo_id=repo_id, filename=filename, cache_dir=str(artifacts_path), local_files_only=True, ) except LocalEntryNotFoundError as err: raise FileNotFoundError( 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}"' ) from err else: checkpoint_path = hf_hub_download( repo_id=repo_id, filename=filename ) self.model = whisper.load_model( name=checkpoint_path, device=self.device ) elif artifacts_path is not None: _log.info(f"loading {self.model_name} from {artifacts_path}") self.model = whisper.load_model( name=self.model_name, device=self.device, download_root=str(artifacts_path), ) else: self.model = whisper.load_model( name=self.model_name, device=self.device ) self.verbose = asr_options.verbose self.timestamps = asr_options.timestamps self.word_timestamps = asr_options.word_timestamps self.language = asr_options.language self.beam_size = asr_options.beam_size self.condition_on_previous_text = asr_options.condition_on_previous_text self.temperature = asr_options.temperature def run(self, conv_res: ConversionResult) -> ConversionResult: # Access the file path from the backend, similar to other pipelines path_or_stream = conv_res.input._backend.path_or_stream # Handle both Path and BytesIO inputs temp_file_path: Path | None = None if isinstance(path_or_stream, BytesIO): # For BytesIO, write to a temporary file (whisper needs a file path) suffix = Path(conv_res.input.file.name).suffix or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file: tmp_file.write(path_or_stream.getvalue()) temp_file_path = Path(tmp_file.name) audio_path = temp_file_path elif isinstance(path_or_stream, Path): audio_path = path_or_stream else: raise RuntimeError( f"ASR pipeline requires a file path or BytesIO stream, " f"but got {type(path_or_stream)}" ) try: if shutil.which("ffmpeg") is None: _log.error(MISSING_FFMPEG_MESSAGE) conv_res.errors.append( ErrorItem( component_type=DoclingComponentType.PIPELINE, module_name="AsrPipeline", error_message=MISSING_FFMPEG_MESSAGE, ) ) conv_res.status = ConversionStatus.FAILURE return conv_res conversation = self.transcribe(audio_path) _process_conversation(conversation, conv_res) return conv_res except Exception as exc: _log.error(f"Audio transcription has an error: {exc}") conv_res.status = ConversionStatus.FAILURE return conv_res finally: # Clean up temporary file if created if temp_file_path is not None and temp_file_path.exists(): try: temp_file_path.unlink() except Exception as e: _log.warning( f"Failed to delete temporary file {temp_file_path}: {e}" ) def transcribe(self, fpath: Path) -> list[_ConversationItem]: result = self.model.transcribe( str(fpath), verbose=self.verbose, language=self.language, word_timestamps=self.word_timestamps, beam_size=self.beam_size, condition_on_previous_text=self.condition_on_previous_text, temperature=self.temperature, ) convo: list[_ConversationItem] = [] for _ in result["segments"]: item = _ConversationItem( start_time=_["start"], end_time=_["end"], text=_["text"], words=[] ) if "words" in _ and self.word_timestamps: item.words = [] for __ in _["words"]: item.words.append( _ConversationWord( start_time=__["start"], end_time=__["end"], text=__["word"], ) ) convo.append(item) return convo class _MlxWhisperModel: def __init__( self, enabled: bool, artifacts_path: Path | None, accelerator_options: AcceleratorOptions, asr_options: InlineAsrMlxWhisperOptions, ): """Transcriber using MLX Whisper for Apple Silicon optimization.""" self.enabled = enabled _log.info(f"artifacts-path: {artifacts_path}") _log.info(f"accelerator_options: {accelerator_options}") if self.enabled: try: import mlx_whisper # type: ignore except ImportError: raise ImportError( "mlx-whisper is not installed. Please install it via " "`pip install mlx-whisper` or do `uv sync --extra asr`." ) self.asr_options = asr_options self.mlx_whisper = mlx_whisper self.device = decide_device( accelerator_options.device, supported_devices=asr_options.supported_devices, ) _log.info(f"Available device for MLX Whisper: {self.device}") self.model_name = asr_options.repo_id _log.info(f"loading _MlxWhisperModel({self.model_name})") # MLX Whisper models are loaded differently - they use HuggingFace repos self.model_path = self.model_name # Store MLX-specific options self.language = asr_options.language self.task = asr_options.task self.word_timestamps = asr_options.word_timestamps self.no_speech_threshold = asr_options.no_speech_threshold self.logprob_threshold = asr_options.logprob_threshold self.compression_ratio_threshold = asr_options.compression_ratio_threshold def run(self, conv_res: ConversionResult) -> ConversionResult: path_or_stream = conv_res.input._backend.path_or_stream temp_file_path: Path | None = None if not isinstance(path_or_stream, (BytesIO, Path)): raise RuntimeError( f"ASR pipeline requires a file path or BytesIO stream, " f"but got {type(path_or_stream)}" ) try: if isinstance(path_or_stream, BytesIO): suffix = Path(conv_res.input.file.name).suffix or ".wav" with tempfile.NamedTemporaryFile( delete=False, suffix=suffix ) as tmp_file: temp_file_path = Path(tmp_file.name) tmp_file.write(path_or_stream.getvalue()) audio_path = temp_file_path else: audio_path = path_or_stream if shutil.which("ffmpeg") is None: _log.error(MISSING_FFMPEG_MESSAGE) conv_res.errors.append( ErrorItem( component_type=DoclingComponentType.PIPELINE, module_name="AsrPipeline", error_message=MISSING_FFMPEG_MESSAGE, ) ) conv_res.status = ConversionStatus.FAILURE return conv_res conversation = self.transcribe(audio_path) _process_conversation(conversation, conv_res) conv_res.status = ConversionStatus.SUCCESS return conv_res except Exception as exc: _log.error(f"MLX Audio transcription has an error: {exc}") conv_res.status = ConversionStatus.FAILURE return conv_res finally: if temp_file_path is not None and temp_file_path.exists(): try: temp_file_path.unlink() except Exception as exc: _log.warning( f"Failed to delete temporary file {temp_file_path}: {exc}" ) def transcribe(self, fpath: Path) -> list[_ConversationItem]: """Transcribe audio using MLX Whisper. Args: fpath: Path to audio file Returns: List of conversation items with timestamps """ result = self.mlx_whisper.transcribe( str(fpath), path_or_hf_repo=self.model_path, language=self.language, task=self.task, word_timestamps=self.word_timestamps, no_speech_threshold=self.no_speech_threshold, logprob_threshold=self.logprob_threshold, compression_ratio_threshold=self.compression_ratio_threshold, ) convo: list[_ConversationItem] = [] # MLX Whisper returns segments similar to native Whisper for segment in result.get("segments", []): item = _ConversationItem( start_time=segment.get("start"), end_time=segment.get("end"), text=segment.get("text", "").strip(), words=[], ) # Add word-level timestamps if available if self.word_timestamps and "words" in segment: item.words = [] for word_data in segment["words"]: item.words.append( _ConversationWord( start_time=word_data.get("start"), end_time=word_data.get("end"), text=word_data.get("word", ""), ) ) convo.append(item) return convo class _WhisperS2TModel: """Transcriber using WhisperS2T with CTranslate2 backend for high-speed inference.""" def __init__( self, enabled: bool, artifacts_path: Path | None, accelerator_options: AcceleratorOptions, asr_options: InlineAsrWhisperS2TOptions, ): self.enabled = enabled _log.info(f"artifacts-path: {artifacts_path}") _log.info(f"accelerator_options: {accelerator_options}") if self.enabled: try: import whisper_s2t # type: ignore except ImportError: raise ImportError( "whisper_s2t is not installed. Please install it via " "`pip install 'whisper-s2t-reborn[pyav]>=1.7.1'`." ) self.whisper_s2t = whisper_s2t self.asr_options = asr_options raw_device = decide_device( accelerator_options.device, supported_devices=asr_options.supported_devices, ) self.device, self.device_index = self._parse_device(raw_device) _log.info( f"Available device for WhisperS2T: {self.device} (index: {self.device_index})" ) self.model_identifier = asr_options.repo_id _log.info(f"loading _WhisperS2TModel({self.model_identifier})") # CTranslate2 does not support float16 or bfloat16 for CPU # 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"): _log.warning( f"compute_type='{compute_type}' is not supported by " f"CTranslate2 on CPU; falling back to 'float32'." ) compute_type = "float32" # Build ASR options for whisper_s2t asr_opts = { "beam_size": asr_options.beam_size, "word_timestamps": asr_options.word_timestamps, } # Build model kwargs model_kwargs = { "device": self.device, "device_index": self.device_index, "compute_type": compute_type, "cpu_threads": asr_options.num_threads, "asr_options": asr_opts, } # large-v3, distil-large-v3, distil-large-v3.5, and large-v3-turbo models require n_mels=128 if self.model_identifier in [ "large-v3", "distil-large-v3", "distil-large-v3.5", "large-v3-turbo", ]: model_kwargs["n_mels"] = 128 self.model = whisper_s2t.load_model( model_identifier=self.model_identifier, **model_kwargs, ) # Store options for transcription self.language = asr_options.language self.task = asr_options.task self.batch_size = asr_options.batch_size self.initial_prompt = asr_options.initial_prompt self.word_timestamps = asr_options.word_timestamps def _parse_device(self, device_str: str) -> tuple: """Parse device string like 'cuda:0' into ('cuda', 0).""" if ":" in device_str: parts = device_str.split(":") device = parts[0] try: device_index = int(parts[1]) except (ValueError, IndexError): device_index = 0 return device, device_index return device_str, 0 def run(self, conv_res: ConversionResult) -> ConversionResult: path_or_stream = conv_res.input._backend.path_or_stream temp_file_path: Path | None = None if isinstance(path_or_stream, BytesIO): suffix = Path(conv_res.input.file.name).suffix or ".wav" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file: tmp_file.write(path_or_stream.getvalue()) temp_file_path = Path(tmp_file.name) audio_path = temp_file_path elif isinstance(path_or_stream, Path): audio_path = path_or_stream else: raise RuntimeError( f"ASR pipeline requires a file path or BytesIO stream, " f"but got {type(path_or_stream)}" ) try: conversation = self.transcribe(audio_path) _process_conversation(conversation, conv_res) conv_res.status = ConversionStatus.SUCCESS return conv_res except Exception as exc: _log.error(f"WhisperS2T transcription error: {exc}") conv_res.status = ConversionStatus.FAILURE return conv_res finally: if temp_file_path is not None and temp_file_path.exists(): try: temp_file_path.unlink() except Exception as e: _log.warning( f"Failed to delete temporary file {temp_file_path}: {e}" ) def transcribe(self, fpath: Path) -> list[_ConversationItem]: """ Transcribe audio using WhisperS2T. Args: fpath: Path to audio file Returns: List of conversation items with timestamps """ out = self.model.transcribe_with_vad( [str(fpath)], lang_codes=[self.language], tasks=[self.task], initial_prompts=[self.initial_prompt], batch_size=self.batch_size, ) convo: list[_ConversationItem] = [] if out and len(out) > 0: for segment in out[0]: words = [] if self.word_timestamps and "word_timestamps" in segment: for w in segment["word_timestamps"]: words.append( _ConversationWord( start_time=w.get("start"), end_time=w.get("end"), text=w.get("word", ""), ) ) item = _ConversationItem( start_time=segment.get("start_time"), end_time=segment.get("end_time"), text=segment.get("text", "").strip(), words=words if words else None, ) convo.append(item) return convo # ============================================================ # Backward-compatibility aliases (private names retained so existing # imports keep working). # ============================================================ # ============================================================ # Transcriber protocol + factory (new; consumed by AsrPipeline and VideoPipeline) # ============================================================ def _merge_into_sentences( items: list["_ConversationItem"], ) -> list["_ConversationItem"]: """Merge Whisper segments into complete sentences. Consecutive segments are merged until a sentence-ending punctuation mark (. ? !) is found. The merged item spans the full time range of all contributing segments and concatenates their text. This produces one block per sentence, which maps cleanly to one speaker per block for diarization. Args: items: ASR segments in chronological order. Returns: One merged item per sentence. """ if not items: return [] merged: list[_ConversationItem] = [] current: _ConversationItem | None = None for item in items: if current is None: current = _ConversationItem( start_time=item.start_time, end_time=item.end_time, text=item.text.strip(), speaker=item.speaker, words=list(item.words or []), ) else: current.end_time = item.end_time current.text = current.text.rstrip() + " " + item.text.strip() if item.words: current.words = (current.words or []) + list(item.words) # Flush on sentence boundary if current.text.rstrip().endswith((".", "?", "!")): merged.append(current) current = None # Flush any remaining text if current is not None: merged.append(current) return merged class _AsrTranscriber(Protocol): """Structural type for ASR backends. Any object exposing `run(conv_res) -> ConversionResult` and `transcribe(fpath: Path) -> list[_ConversationItem]` satisfies this. """ def transcribe(self, fpath: Path) -> list[_ConversationItem]: ... def run(self, conv_res: ConversionResult) -> ConversionResult: ... class _AsrModelFactory: """Builds a concrete ASR transcriber from ASR options. Shared by AsrPipeline and (later) VideoPipeline so both construct backends the same way instead of duplicating the isinstance chain. """ @staticmethod def create( asr_options: InlineAsrOptions, artifacts_path: Path | None, accelerator_options: AcceleratorOptions, ) -> _AsrTranscriber: if isinstance(asr_options, InlineAsrNativeWhisperOptions): return _NativeWhisperModel( enabled=True, artifacts_path=artifacts_path, accelerator_options=accelerator_options, asr_options=asr_options, ) elif isinstance(asr_options, InlineAsrMlxWhisperOptions): return _MlxWhisperModel( enabled=True, artifacts_path=artifacts_path, accelerator_options=accelerator_options, asr_options=asr_options, ) elif isinstance(asr_options, InlineAsrWhisperS2TOptions): return _WhisperS2TModel( enabled=True, artifacts_path=artifacts_path, accelerator_options=accelerator_options, asr_options=asr_options, ) raise ValueError(f"No ASR model support for {asr_options}")