chore: import upstream snapshot with attribution

unstructured/partition/utils/speech_to_text/__init__.py

@@ -0,0 +1,8 @@
+"""Speech-to-text agents for transcribing audio in the multimodal partition pipeline."""
+
+from unstructured.partition.utils.speech_to_text.speech_to_text_interface import (
+    SpeechToTextAgent,
+    TranscriptionSegment,
+)
+
+__all__ = ["SpeechToTextAgent", "TranscriptionSegment"]

unstructured/partition/utils/speech_to_text/speech_to_text_interface.py

@@ -0,0 +1,115 @@
+"""Abstract interface for speech-to-text (STT) agents used by the audio partitioner."""
+
+from __future__ import annotations
+
+import functools
+import importlib
+from abc import ABC, abstractmethod
+from typing import TypedDict
+
+from unstructured.logger import logger
+from unstructured.partition.utils.config import env_config
+from unstructured.partition.utils.constants import STT_AGENT_MODULES_WHITELIST
+
+
+class TranscriptionSegment(TypedDict):
+    """A single segment of a transcription with text and timestamps in seconds.
+
+    The audio partitioner (caller of :meth:`SpeechToTextAgent.transcribe_segments`) treats
+    empty or whitespace-only ``text`` as absent: it strips each segment's text and skips
+    segments that are empty after stripping. Agents may therefore return such segments
+    without filtering; the partitioner is the single place that drops them.
+    """
+
+    text: str
+    start: float
+    end: float
+
+
+class SpeechToTextAgent(ABC):
+    """Defines the interface for a speech-to-text transcription service."""
+
+    @classmethod
+    def get_agent(cls, agent_module: str | None = None) -> "SpeechToTextAgent":
+        """Return the configured SpeechToTextAgent instance.
+
+        The agent module is resolved from `agent_module` when provided, otherwise from
+        the `STT_AGENT` environment variable (default: Whisper).
+        """
+        return cls.get_instance(agent_module or env_config.STT_AGENT)
+
+    @staticmethod
+    @functools.lru_cache(maxsize=env_config.STT_AGENT_CACHE_SIZE)
+    def get_instance(agent_module: str) -> "SpeechToTextAgent":
+        """Load and return a SpeechToTextAgent for the given fully-qualified class name.
+
+        Results are cached (keyed on ``agent_module``) up to ``STT_AGENT_CACHE_SIZE``
+        entries. Because the cache key is the class name only, model-configuration
+        environment variables (``WHISPER_MODEL_SIZE``, ``WHISPER_DEVICE``,
+        ``WHISPER_FP16``) are read once at first instantiation and ignored on subsequent
+        calls. A process restart is required to pick up configuration changes.
+        """
+        module_name, class_name = agent_module.rsplit(".", 1)
+        if module_name not in STT_AGENT_MODULES_WHITELIST:
+            raise ValueError(
+                f"Speech-to-text agent module {module_name} must be in the whitelist: "
+                f"{STT_AGENT_MODULES_WHITELIST}."
+            )
+        try:
+            mod = importlib.import_module(module_name)
+            loaded_class = getattr(mod, class_name)
+        except (ImportError, AttributeError) as e:
+            logger.error(f"Failed to load SpeechToTextAgent class '{agent_module}': {e}")
+            raise RuntimeError(
+                f"Could not load the SpeechToText agent class '{agent_module}'. "
+                "Install the audio extra: "
+                'pip install "unstructured[audio]"'
+            ) from e
+
+        if not isinstance(loaded_class, type):
+            raise TypeError(
+                f"'{agent_module}' does not refer to a class "
+                f"(got {type(loaded_class).__name__}). "
+                "Speech-to-text agent must be a subclass of SpeechToTextAgent."
+            )
+        if not issubclass(loaded_class, SpeechToTextAgent):
+            raise TypeError(
+                f"'{agent_module}' must be a subclass of SpeechToTextAgent, "
+                f"got {loaded_class.__qualname__}."
+            )
+        try:
+            return loaded_class()
+        except Exception as e:
+            logger.error(f"SpeechToTextAgent '{class_name}' loaded but failed to initialize: {e}")
+            raise RuntimeError(
+                f"SpeechToText agent '{class_name}' was imported successfully but its "
+                f"constructor raised an error. "
+                f"Original error: {e}"
+            ) from e
+
+    @abstractmethod
+    def transcribe_segments(
+        self, audio_path: str, *, language: str | None = None
+    ) -> list[TranscriptionSegment]:
+        """Transcribe audio and return segment-level results with timestamps.
+
+        This is the **primary method** to implement. All partitioning calls go through
+        here. Subclasses that support segment-level output (e.g. Whisper) should return
+        one entry per segment; subclasses without native segment support should return a
+        single segment with ``start=0.0`` and ``end=0.0``.
+
+        Parameters
+        ----------
+        audio_path
+            Path to an audio file (e.g. WAV, MP3).
+        language
+            Optional ISO 639-1 language code for the spoken language (e.g. ``"en"``).
+            When ``None``, the agent may auto-detect.
+
+        Returns
+        -------
+        List of :class:`TranscriptionSegment` dicts, each with ``"text"``, ``"start"``,
+        and ``"end"`` keys. Return an empty list when the audio contains no speech.
+        Segments with empty or whitespace-only ``text`` may be included; the partitioner
+        will strip and drop them.
+        """

unstructured/partition/utils/speech_to_text/whisper_stt.py

@@ -0,0 +1,85 @@
+"""Whisper-based speech-to-text agent for the audio partitioner."""
+
+from __future__ import annotations
+
+import threading
+
+from unstructured.partition.utils.config import env_config
+from unstructured.partition.utils.speech_to_text.speech_to_text_interface import (
+    SpeechToTextAgent,
+    TranscriptionSegment,
+)
+
+
+class SpeechToTextAgentWhisper(SpeechToTextAgent):
+    """Speech-to-text implementation using OpenAI Whisper.
+
+    **Concurrency model** — a single instance is shared across all callers via the
+    ``lru_cache`` in :meth:`~SpeechToTextAgent.get_instance`. Because
+    ``whisper.model.transcribe()`` is not documented as thread-safe, a per-instance
+    ``threading.Lock`` serializes all transcription calls. This prevents data races but
+    means the process can only run one transcription at a time for the default agent —
+    a hidden throughput ceiling under concurrent workloads. For true parallelism, use
+    process-based concurrency (e.g. ``multiprocessing`` or separate worker processes)
+    rather than threads.
+
+    **Configuration snapshot** — model size, device, and FP16 flag are read from
+    environment variables (``WHISPER_MODEL_SIZE``, ``WHISPER_DEVICE``, ``WHISPER_FP16``)
+    at construction time and frozen for the lifetime of the cached instance. Changing
+    those variables after the first call has no effect without a process restart.
+    """
+
+    def __init__(self, model_size: str | None = None) -> None:
+        """Initialize the Whisper model.
+
+        Parameters
+        ----------
+        model_size
+            Whisper model size: "tiny", "base", "small", "medium", "large", or "large-v3".
+            Larger models are more accurate but slower and use more memory.
+            When None, uses the WHISPER_MODEL_SIZE environment variable (default \"base\").
+        """
+        import whisper
+
+        size = model_size if model_size is not None else env_config.WHISPER_MODEL_SIZE
+        device = env_config.WHISPER_DEVICE.strip() or None  # empty -> auto
+        try:
+            self._model = whisper.load_model(size, device=device)
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to load Whisper model '{size}' on device {device!r}. "
+                "Possible causes: invalid model name, network error during model download, "
+                "or insufficient GPU memory (CUDA OOM). "
+                f"Valid model sizes: tiny, base, small, medium, large, large-v3. "
+                f"Original error: {exc}"
+            ) from exc
+        self._fp16 = env_config.WHISPER_FP16
+        self._lock = threading.Lock()
+
+    def transcribe_segments(
+        self, audio_path: str, *, language: str | None = None
+    ) -> list[TranscriptionSegment]:
+        """Transcribe audio and return one segment per Whisper segment with timestamps."""
+        options: dict = {"fp16": self._fp16}
+        if language is not None:
+            options["language"] = language
+        with self._lock:
+            result = self._model.transcribe(audio_path, **options)
+        segments: list[TranscriptionSegment] = []
+        for seg in result.get("segments", []):
+            segments.append(
+                TranscriptionSegment(
+                    text=(seg.get("text") or ""),
+                    start=float(seg.get("start", 0)),
+                    end=float(seg.get("end", 0)),
+                )
+            )
+        if not segments and result.get("text", "").strip():
+            segments = [
+                TranscriptionSegment(
+                    text=result.get("text", "").strip(),
+                    start=0.0,
+                    end=0.0,
+                )
+            ]
+        return segments