docling 2.69.0__py3-none-any.whl

docling/pipeline/asr_pipeline.py

@@ -0,0 +1,431 @@
+import logging
+import os
+import re
+import sys
+import tempfile
+from io import BytesIO
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional, Union, cast
+
+from docling_core.types.doc import DoclingDocument, DocumentOrigin
+
+# import whisper  # type: ignore
+# import librosa
+# import numpy as np
+# import soundfile as sf  # type: ignore
+from docling_core.types.doc.labels import DocItemLabel
+from pydantic import BaseModel, Field, validator
+
+from docling.backend.abstract_backend import AbstractDocumentBackend
+from docling.backend.noop_backend import NoOpBackend
+
+# from pydub import AudioSegment  # type: ignore
+# from transformers import WhisperForConditionalGeneration, WhisperProcessor, pipeline
+from docling.datamodel.accelerator_options import (
+    AcceleratorOptions,
+)
+from docling.datamodel.base_models import (
+    ConversionStatus,
+    FormatToMimeType,
+)
+from docling.datamodel.document import ConversionResult, InputDocument
+from docling.datamodel.pipeline_options import (
+    AsrPipelineOptions,
+)
+from docling.datamodel.pipeline_options_asr_model import (
+    InlineAsrMlxWhisperOptions,
+    InlineAsrNativeWhisperOptions,
+    # AsrResponseFormat,
+    InlineAsrOptions,
+)
+from docling.datamodel.pipeline_options_vlm_model import (
+    InferenceFramework,
+)
+from docling.datamodel.settings import settings
+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__)
+
+
+class _ConversationWord(BaseModel):
+    text: str
+    start_time: Optional[float] = Field(
+        None, description="Start time in seconds from video start"
+    )
+    end_time: Optional[float] = Field(
+        None, ge=0, description="End time in seconds from video start"
+    )
+
+
+class _ConversationItem(BaseModel):
+    text: str
+    start_time: Optional[float] = Field(
+        None, description="Start time in seconds from video start"
+    )
+    end_time: Optional[float] = Field(
+        None, ge=0, description="End time in seconds from video start"
+    )
+    speaker_id: Optional[int] = Field(None, description="Numeric speaker identifier")
+    speaker: Optional[str] = Field(
+        None, description="Speaker name, defaults to speaker-{speaker_id}"
+    )
+    words: Optional[list[_ConversationWord]] = 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
+
+
+class _NativeWhisperModel:
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        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.temperature = asr_options.temperature
+
+            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})")
+            if 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
+
+    def run(self, conv_res: ConversionResult) -> ConversionResult:
+        # Access the file path from the backend, similar to how other pipelines handle it
+        path_or_stream = conv_res.input._backend.path_or_stream
+
+        # Handle both Path and BytesIO inputs
+        temp_file_path: Optional[Path] = None
+
+        if isinstance(path_or_stream, BytesIO):
+            # For BytesIO, write to a temporary file since whisper requires 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, but got {type(path_or_stream)}"
+            )
+
+        try:
+            conversation = self.transcribe(audio_path)
+
+            # Ensure we have a proper DoclingDocument
+            origin = DocumentOrigin(
+                filename=conv_res.input.file.name or "audio.wav",
+                mimetype="audio/x-wav",
+                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:
+                conv_res.document.add_text(
+                    label=DocItemLabel.TEXT, text=citem.to_string()
+                )
+
+            return conv_res
+
+        except Exception as exc:
+            _log.error(f"Audio tranciption 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, word_timestamps=self.word_timestamps
+        )
+
+        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: Optional[Path],
+        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:
+        audio_path: Path = Path(conv_res.input.file).resolve()
+
+        try:
+            conversation = self.transcribe(audio_path)
+
+            # Ensure we have a proper DoclingDocument
+            origin = DocumentOrigin(
+                filename=conv_res.input.file.name or "audio.wav",
+                mimetype="audio/x-wav",
+                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:
+                conv_res.document.add_text(
+                    label=DocItemLabel.TEXT, text=citem.to_string()
+                )
+
+            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
+
+    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 AsrPipeline(BasePipeline):
+    def __init__(self, pipeline_options: AsrPipelineOptions):
+        super().__init__(pipeline_options)
+        self.keep_backend = True
+
+        self.pipeline_options: AsrPipelineOptions = pipeline_options
+        self._model: Union[_NativeWhisperModel, _MlxWhisperModel]
+
+        if isinstance(self.pipeline_options.asr_options, InlineAsrNativeWhisperOptions):
+            native_asr_options: InlineAsrNativeWhisperOptions = (
+                self.pipeline_options.asr_options
+            )
+            self._model = _NativeWhisperModel(
+                enabled=True,  # must be always enabled for this pipeline to make sense.
+                artifacts_path=self.artifacts_path,
+                accelerator_options=pipeline_options.accelerator_options,
+                asr_options=native_asr_options,
+            )
+        elif isinstance(self.pipeline_options.asr_options, InlineAsrMlxWhisperOptions):
+            mlx_asr_options: InlineAsrMlxWhisperOptions = (
+                self.pipeline_options.asr_options
+            )
+            self._model = _MlxWhisperModel(
+                enabled=True,  # must be always enabled for this pipeline to make sense.
+                artifacts_path=self.artifacts_path,
+                accelerator_options=pipeline_options.accelerator_options,
+                asr_options=mlx_asr_options,
+            )
+        else:
+            _log.error(f"No model support for {self.pipeline_options.asr_options}")
+
+    def _has_text(self, document: "DoclingDocument") -> bool:
+        """
+        Helper method to check if the document contains any transcribed text.
+        A transcription is considered non-empty if the .texts list contains items with actual, non whitespace content.
+        """
+        if not document or not document.texts:
+            return False
+        for item in document.texts:
+            if item.text and item.text.strip():
+                return True
+        return False
+
+    def _determine_status(self, conv_res: ConversionResult) -> ConversionStatus:
+        """Determines the final status of ASR Conversion based on its result."""
+        if conv_res.status == ConversionStatus.FAILURE or conv_res.errors:
+            return ConversionStatus.FAILURE
+        if not self._has_text(conv_res.document):
+            _log.warning(
+                "ASR conversion resulted in an empty document."
+                f"File: {conv_res.input.file.name}"
+            )
+            return ConversionStatus.PARTIAL_SUCCESS
+        return ConversionStatus.SUCCESS
+
+    @classmethod
+    def get_default_options(cls) -> AsrPipelineOptions:
+        return AsrPipelineOptions()
+
+    def _build_document(self, conv_res: ConversionResult) -> ConversionResult:
+        _log.info(f"start _build_document in AsrPipeline: {conv_res.input.file}")
+        with TimeRecorder(conv_res, "doc_build", scope=ProfilingScope.DOCUMENT):
+            self._model.run(conv_res=conv_res)
+
+        return conv_res
+
+    @classmethod
+    def is_backend_supported(cls, backend: AbstractDocumentBackend):
+        return isinstance(backend, NoOpBackend)

docling/pipeline/base_extraction_pipeline.py

@@ -0,0 +1,72 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+from docling.datamodel.base_models import ConversionStatus, ErrorItem
+from docling.datamodel.document import InputDocument
+from docling.datamodel.extraction import ExtractionResult, ExtractionTemplateType
+from docling.datamodel.pipeline_options import BaseOptions, PipelineOptions
+from docling.datamodel.settings import settings
+
+_log = logging.getLogger(__name__)
+
+
+class BaseExtractionPipeline(ABC):
+    def __init__(self, pipeline_options: PipelineOptions):
+        self.pipeline_options = pipeline_options
+
+        self.artifacts_path: Optional[Path] = None
+        if pipeline_options.artifacts_path is not None:
+            self.artifacts_path = Path(pipeline_options.artifacts_path).expanduser()
+        elif settings.artifacts_path is not None:
+            self.artifacts_path = Path(settings.artifacts_path).expanduser()
+
+        if self.artifacts_path is not None and not self.artifacts_path.is_dir():
+            raise RuntimeError(
+                f"The value of {self.artifacts_path=} is not valid. "
+                "When defined, it must point to a folder containing all models required by the pipeline."
+            )
+
+    def execute(
+        self,
+        in_doc: InputDocument,
+        raises_on_error: bool,
+        template: Optional[ExtractionTemplateType] = None,
+    ) -> ExtractionResult:
+        ext_res = ExtractionResult(input=in_doc)
+
+        try:
+            ext_res = self._extract_data(ext_res, template)
+            ext_res.status = self._determine_status(ext_res)
+        except Exception as e:
+            ext_res.status = ConversionStatus.FAILURE
+            error_item = ErrorItem(
+                component_type="extraction_pipeline",
+                module_name=self.__class__.__name__,
+                error_message=str(e),
+            )
+            ext_res.errors.append(error_item)
+            if raises_on_error:
+                raise e
+
+        return ext_res
+
+    @abstractmethod
+    def _extract_data(
+        self,
+        ext_res: ExtractionResult,
+        template: Optional[ExtractionTemplateType] = None,
+    ) -> ExtractionResult:
+        """Subclass must populate ext_res.pages/errors and return the result."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _determine_status(self, ext_res: ExtractionResult) -> ConversionStatus:
+        """Subclass must decide SUCCESS/PARTIAL_SUCCESS/FAILURE based on ext_res."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_default_options(cls) -> PipelineOptions:
+        pass