textflowkit 0.1.3__py3-none-any.whl

textflowkit/core/diarize.py

@@ -0,0 +1,229 @@
+"""Speaker diarization.
+
+The `Segment.speaker` field existed before anything could fill it, which is the
+same "surface with no implementation" defect the audit caught on
+`--speaker-labels`. This module fills it - and refuses to pretend:
+
+- diarization is **opt-in**
+- if the backend is missing or unconfigured, the job **fails with an actionable
+  error**; it never silently returns a transcript with empty speakers
+- the real backend (`pyannote.audio`) is an optional extra and its pretrained
+  pipeline is gated behind a Hugging Face token, so the live path is not covered
+  by CI. The *assignment* logic is fully tested with a stub.
+
+Design: a diarizer answers "who spoke when" as a list of turns over the audio.
+Assigning speakers to transcript segments is a separate, pure step, so it can be
+tested without any model at all.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+from typing import Protocol
+
+from textflowkit.core.engine import _pick_device
+from textflowkit.core.model import Segment
+
+ENV_HF_TOKEN = "HF_TOKEN"
+ENV_PYANNOTE_MODEL = "TEXTFLOWKIT_PYANNOTE_MODEL"
+ENV_DIARIZE_DEVICE = "TEXTFLOWKIT_DIARIZE_DEVICE"
+DEFAULT_PYANNOTE_MODEL = "pyannote/speaker-diarization-3.1"
+
+
+class DiarizationError(RuntimeError):
+    """Raised when diarization was requested but could not be performed."""
+
+
+@dataclass(slots=True)
+class SpeakerTurn:
+    """A span of audio attributed to one speaker."""
+
+    start: float
+    end: float
+    speaker: str
+
+    @property
+    def duration(self) -> float:
+        return max(0.0, self.end - self.start)
+
+
+class Diarizer(Protocol):
+    """Identifies who spoke when."""
+
+    name: str
+
+    def diarize(self, audio_path: str | Path) -> list[SpeakerTurn]: ...
+
+
+def overlap(a_start: float, a_end: float, b_start: float, b_end: float) -> float:
+    """Seconds of overlap between two spans. Zero when they do not intersect."""
+    return max(0.0, min(a_end, b_end) - max(a_start, b_start))
+
+
+def assign_speakers(
+    segments: list[Segment],
+    turns: list[SpeakerTurn],
+    *,
+    min_overlap: float = 0.0,
+) -> int:
+    """Label each segment with the speaker who overlaps it most.
+
+    A segment is assigned to the turn with the greatest time overlap. Ties go to
+    the earlier turn, so the result is deterministic rather than dependent on
+    iteration order. A segment with no overlapping turn is left unlabelled
+    rather than guessed at.
+
+    Returns the number of segments that received a label.
+    """
+    if not turns:
+        return 0
+
+    labelled = 0
+    for segment in segments:
+        best: SpeakerTurn | None = None
+        best_overlap = 0.0
+        for turn in turns:
+            amount = overlap(segment.start, segment.end, turn.start, turn.end)
+            # Strictly greater wins. On an exact tie the earlier turn wins, so
+            # the result depends on the turns themselves and not on the order
+            # they happen to arrive in.
+            if amount > best_overlap or (
+                amount == best_overlap and best is not None and turn.start < best.start
+            ):
+                best = turn
+                best_overlap = amount
+        if best is not None and best_overlap > min_overlap:
+            segment.speaker = best.speaker
+            labelled += 1
+    return labelled
+
+
+class PyannoteDiarizer:
+    """Diarization via `pyannote.audio`.
+
+    Optional dependency, and the pretrained pipeline is gated: it needs a Hugging
+    Face token with access granted to the model. Both conditions are checked up
+    front so the failure names what is missing instead of surfacing as an opaque
+    load error.
+    """
+
+    name = "pyannote"
+
+    def __init__(
+        self, model: str | None = None, token: str | None = None,
+        device: str | None = None,
+    ) -> None:
+        self.model_name = model or os.environ.get(ENV_PYANNOTE_MODEL, DEFAULT_PYANNOTE_MODEL)
+        self._token = token or os.environ.get(ENV_HF_TOKEN)
+        self.device = _pick_device(device or os.environ.get(ENV_DIARIZE_DEVICE))
+        self._pipeline = None
+        self._lock = threading.RLock()
+
+    def _load(self):
+        if self._pipeline is not None:
+            return self._pipeline
+        try:
+            from pyannote.audio import Pipeline
+        except ImportError as exc:
+            raise DiarizationError(
+                "diarization requires the optional 'diarize' extra. "
+                "Install with: pip install 'textflowkit[diarize]'"
+            ) from exc
+        if not self._token:
+            raise DiarizationError(
+                f"diarization requires a Hugging Face token with access to "
+                f"'{self.model_name}'. Set {ENV_HF_TOKEN}, or pass token=... . "
+                "The model is gated, so access must also be granted on Hugging Face."
+            )
+        # pyannote.audio renamed `use_auth_token` to `token` in 4.0. Pass the
+        # keyword the installed version actually accepts rather than pinning the
+        # code to one release.
+        try:
+            import inspect
+
+            params = inspect.signature(Pipeline.from_pretrained).parameters
+            kwargs = {"token": self._token} if "token" in params else {"use_auth_token": self._token}
+            loaded = Pipeline.from_pretrained(self.model_name, **kwargs)
+            if hasattr(loaded, "to"):
+                import torch
+
+                loaded.to(torch.device(self.device))
+            self._pipeline = loaded
+        except Exception as exc:  # provider errors vary
+            raise DiarizationError(
+                f"could not load diarization model '{self.model_name}': {exc}"
+            ) from exc
+        return self._pipeline
+
+    @staticmethod
+    def _load_waveform(audio_path: str | Path) -> dict:
+        """Read audio ourselves instead of letting pyannote decode it.
+
+        pyannote.audio 4.x decodes through `torchcodec`, whose bundled DLLs are
+        built against specific torch releases and fail to load against the ROCm
+        torch build this project targets. Its own error message names the way
+        out: "provide audio as a waveform dictionary". Reading with `soundfile`
+        (already a pyannote dependency) avoids that native-extension coupling
+        entirely and skips a decode step.
+        """
+        try:
+            import soundfile as sf
+            import torch
+        except ImportError as exc:  # pragma: no cover - both are hard deps of pyannote
+            raise DiarizationError(f"diarization requires soundfile and torch: {exc}") from exc
+
+        try:
+            data, sample_rate = sf.read(str(audio_path), dtype="float32", always_2d=True)
+        except Exception as exc:
+            raise DiarizationError(f"could not read audio '{audio_path}': {exc}") from exc
+
+        waveform = torch.from_numpy(data.T)  # (channels, samples)
+        if waveform.shape[0] > 1:
+            # pyannote expects mono; averaging is what it documents for
+            # multi-channel input.
+            waveform = waveform.mean(dim=0, keepdim=True)
+        return {"waveform": waveform, "sample_rate": sample_rate}
+
+    def diarize(self, audio_path: str | Path) -> list[SpeakerTurn]:
+        try:
+            with self._lock:
+                pipeline = self._load()
+                annotation = pipeline(self._load_waveform(audio_path))
+        except DiarizationError:
+            raise
+        except Exception as exc:
+            raise DiarizationError(f"diarization failed: {exc}") from exc
+
+        # pyannote.audio 4.x returns a DiarizeOutput wrapper exposing the
+        # annotation as `.speaker_diarization`; earlier releases returned the
+        # Annotation directly. Accept either rather than pinning to one version.
+        annotation = getattr(annotation, "speaker_diarization", annotation)
+
+        turns: list[SpeakerTurn] = []
+        for turn, _, speaker in annotation.itertracks(yield_label=True):
+            turns.append(SpeakerTurn(start=turn.start, end=turn.end, speaker=str(speaker)))
+        return turns
+
+
+_DIARIZER_CACHE_LOCK = threading.Lock()
+
+
+@lru_cache(maxsize=4)
+def _cached_diarizer(model: str, token: str | None, device: str) -> PyannoteDiarizer:
+    return PyannoteDiarizer(model=model, token=token, device=device)
+
+
+def get_diarizer(backend: str = "pyannote", **kwargs) -> Diarizer:
+    if backend in ("pyannote", "default"):
+        model = kwargs.pop("model", None) or os.environ.get(ENV_PYANNOTE_MODEL, DEFAULT_PYANNOTE_MODEL)
+        token = kwargs.pop("token", None) or os.environ.get(ENV_HF_TOKEN)
+        device = _pick_device(kwargs.pop("device", None) or os.environ.get(ENV_DIARIZE_DEVICE))
+        if kwargs:
+            raise TypeError(f"unknown diarizer options: {', '.join(kwargs)}")
+        with _DIARIZER_CACHE_LOCK:
+            return _cached_diarizer(model, token, device)
+    raise DiarizationError(f"unknown diarization backend: {backend}")

textflowkit/core/engine.py

@@ -0,0 +1,127 @@
+"""Speech-to-text engines.
+
+The default engine is openai-whisper on PyTorch. On this project's reference
+hardware (AMD Strix Halo) that means ROCm; on NVIDIA it means CUDA; with neither
+it falls back to CPU. The engine interface is intentionally tiny so alternatives
+can be added without touching the pipeline.
+"""
+
+from __future__ import annotations
+
+import threading
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Protocol
+
+from textflowkit.core.model import Segment, Transcript
+
+
+class Engine(Protocol):
+    name: str
+
+    def transcribe(
+        self,
+        audio_path: str | Path,
+        *,
+        language: str | None = None,
+    ) -> Transcript: ...
+
+
+def _pick_device(prefer: str | None = None) -> str:
+    if prefer:
+        return prefer
+    try:
+        import torch
+    except ImportError:
+        return "cpu"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+class WhisperEngine:
+    """openai-whisper backed engine."""
+
+    name = "openai-whisper"
+
+    def __init__(self, model: str = "small", device: str | None = None, fp16: bool | None = None):
+        self.model_name = model
+        self.device = _pick_device(device)
+        if fp16 is None:
+            fp16 = self.device != "cpu"
+        self.fp16 = fp16
+        self._model: Any = None
+        self._lock = threading.RLock()
+
+    def _load(self):
+        with self._lock:
+            if self._model is None:
+                try:
+                    import whisper
+                except ImportError as exc:
+                    raise RuntimeError(
+                        "openai-whisper is not installed. Install with: pip install openai-whisper"
+                    ) from exc
+                self._model = whisper.load_model(self.model_name, device=self.device)
+            return self._model
+
+    def transcribe(
+        self,
+        audio_path: str | Path,
+        *,
+        language: str | None = None,
+    ) -> Transcript:
+        with self._lock:
+            model = self._load()
+            result = model.transcribe(
+                str(audio_path),
+                language=language,
+                fp16=self.fp16,
+                verbose=False,
+                word_timestamps=True,
+            )
+
+        segments: list[Segment] = []
+        for raw in result.get("segments", []) or []:
+            text = str(raw.get("text", "")).strip()
+            if not text:
+                continue
+            segments.append(
+                Segment(
+                    start=float(raw.get("start", 0.0)),
+                    end=float(raw.get("end", 0.0)),
+                    text=text,
+                    speaker=None,
+                )
+            )
+
+        duration = segments[-1].end if segments else None
+        return Transcript(
+            source=str(audio_path),
+            language=result.get("language"),
+            segments=segments,
+            duration=duration,
+            engine=self.name,
+            metadata={"model": self.model_name, "device": self.device},
+        )
+
+
+_ENGINE_CACHE_LOCK = threading.Lock()
+
+
+@lru_cache(maxsize=8)
+def _cached_whisper(model: str, device: str | None, fp16: bool | None) -> WhisperEngine:
+    return WhisperEngine(model=model, device=device, fp16=fp16)
+
+
+def get_engine(name: str = "whisper", **kwargs: Any) -> Engine:
+    if name in ("whisper", "openai-whisper", "default"):
+        model = kwargs.pop("model", "small")
+        device = kwargs.pop("device", None)
+        fp16 = kwargs.pop("fp16", None)
+        if kwargs:
+            raise TypeError(f"unknown Whisper options: {', '.join(kwargs)}")
+        with _ENGINE_CACHE_LOCK:
+            return _cached_whisper(model, device, fp16)
+    raise ValueError(f"unknown engine: {name}")
+

textflowkit/core/executor.py

@@ -0,0 +1,301 @@
+"""Bounded job executor.
+
+Before this, `submit()` spawned one raw `threading.Thread` per job and forgot it.
+That had two consequences the review called out:
+
+- **No concurrency bound.** N submissions meant N simultaneous Whisper runs
+  competing for the same GPU. On a device Whisper already saturates, parallel
+  jobs do not go faster - they thrash memory.
+- **No handle on running work.** Nothing held a reference to a running job, so
+  "cancel" had nothing to act on.
+
+This module replaces the fire-and-forget thread with a fixed worker pool that
+owns job lifecycle. Cancellation is cooperative: a job checks a token at stage
+boundaries and stops cleanly. A job inside a single long model call cannot be
+interrupted mid-call - that call finishes, then the job stops. Stated plainly
+here because it is a real limit, not an implementation detail.
+"""
+
+from __future__ import annotations
+
+import os
+import queue
+import threading
+from typing import Any
+
+from textflowkit.core.cancel import CancelledError
+from textflowkit.core.jobs import Job, JobState, JobStore, get_default_store
+
+ENV_CONCURRENCY = "TEXTFLOWKIT_MAX_CONCURRENCY"
+DEFAULT_CONCURRENCY = 1
+ENV_MAX_PENDING = "TEXTFLOWKIT_MAX_PENDING_JOBS"
+DEFAULT_MAX_PENDING = 100
+
+
+class QueueFullError(RuntimeError):
+    """The executor cannot accept another pending job right now."""
+
+
+# The signal itself lives in a leaf module so the source layer can re-raise it
+# around broad exception handling. This alias keeps the name used everywhere
+# else in the codebase and in tests.
+JobCancelled = CancelledError
+
+
+class CancelToken:
+    """A cooperative cancellation flag, checked at stage boundaries."""
+
+    __slots__ = ("_cancelled", "_lock")
+
+    def __init__(self) -> None:
+        self._cancelled = False
+        self._lock = threading.Lock()
+
+    def cancel(self) -> None:
+        with self._lock:
+            self._cancelled = True
+
+    @property
+    def cancelled(self) -> bool:
+        with self._lock:
+            return self._cancelled
+
+    def checkpoint(self) -> None:
+        """Raise `JobCancelled` if cancellation was requested."""
+        if self.cancelled:
+            raise JobCancelled()
+
+
+class JobExecutor:
+    """A fixed pool of workers that run jobs from a queue.
+
+    The pool size is the concurrency bound. Defaults to 1 because Whisper
+    saturates a GPU on its own; override with `TEXTFLOWKIT_MAX_CONCURRENCY`.
+    """
+
+    def __init__(
+        self,
+        store: JobStore,
+        *,
+        max_concurrency: int | None = None,
+        max_pending: int | None = None,
+    ) -> None:
+        self._store = store
+        if max_concurrency is None:
+            raw = os.environ.get(ENV_CONCURRENCY)
+            try:
+                max_concurrency = int(raw) if raw else DEFAULT_CONCURRENCY
+            except ValueError:
+                max_concurrency = DEFAULT_CONCURRENCY
+        self._max_concurrency = max(1, max_concurrency)
+
+        if max_pending is None:
+            raw = os.environ.get(ENV_MAX_PENDING)
+            try:
+                max_pending = int(raw) if raw else DEFAULT_MAX_PENDING
+            except ValueError:
+                max_pending = DEFAULT_MAX_PENDING
+        self._max_pending = max(1, max_pending)
+        # This semaphore counts queued jobs, not running workers. It also lets
+        # shutdown enqueue sentinels without deadlocking on a full queue.
+        self._pending_slots = threading.BoundedSemaphore(self._max_pending)
+
+        self._queue: queue.Queue[tuple[str, dict[str, Any]] | None] = queue.Queue()
+        self._workers: list[threading.Thread] = []
+        self._tokens: dict[str, CancelToken] = {}
+        self._lock = threading.RLock()
+        self._started = False
+        self._shutdown = False
+
+    # -- lifecycle ---------------------------------------------------------
+
+    @property
+    def store(self) -> JobStore:
+        """The store this executor writes to."""
+        return self._store
+
+    @property
+    def max_concurrency(self) -> int:
+        return self._max_concurrency
+
+    @property
+    def max_pending(self) -> int:
+        return self._max_pending
+
+    def start(self) -> None:
+        """Start worker threads. Idempotent; called lazily by submit()."""
+
+        def _start_locked() -> None:
+            for i in range(self._max_concurrency):
+                t = threading.Thread(
+                    target=self._worker,
+                    name=f"textflowkit-worker-{i}",
+                    daemon=True,
+                )
+                t.start()
+                self._workers.append(t)
+            self._started = True
+
+        with self._lock:
+            if self._started or self._shutdown:
+                return
+            # A durable store may hold jobs left mid-flight by a previous
+            # process. They have no worker now, so fail them rather than
+            # reporting jobs that can never finish. This assumes one owning
+            # process per store, which is the documented deployment model.
+            self._store.reap_incomplete(
+                reason="interrupted by restart; no worker is running this job"
+            )
+            _start_locked()
+
+    def shutdown(self, *, wait: bool = True, timeout: float = 5.0) -> None:
+        """Stop accepting work and drain the pool."""
+        with self._lock:
+            if self._shutdown:
+                return
+            self._shutdown = True
+        for _ in self._workers:
+            self._queue.put(None)  # sentinel: one per worker
+        if wait:
+            for t in self._workers:
+                t.join(timeout=timeout)
+
+    # -- work --------------------------------------------------------------
+
+    def submit(
+        self, *, source: str, request: dict[str, Any] | None = None, **kwargs: Any
+    ) -> Job:
+        """Queue a job and return it immediately."""
+        self.start()
+        with self._lock:
+            if self._shutdown:
+                raise RuntimeError("job executor is shut down")
+            if not self._pending_slots.acquire(blocking=False):
+                raise QueueFullError(
+                    f"job queue is full ({self._max_pending} pending); retry later"
+                )
+            try:
+                job = self._store.create(source, request=request)
+                self._queue.put_nowait((job.id, {"source": source, **kwargs}))
+            except Exception:
+                self._pending_slots.release()
+                raise
+            return job
+
+    def enqueue(self, job: Job, *, source: str, **kwargs: Any) -> Job:
+        """Queue an existing prepared job (the durable resume path)."""
+        self.start()
+        with self._lock:
+            if self._shutdown:
+                raise RuntimeError("job executor is shut down")
+            if not self._pending_slots.acquire(blocking=False):
+                raise QueueFullError(
+                    f"job queue is full ({self._max_pending} pending); retry later"
+                )
+            try:
+                self._queue.put_nowait((job.id, {"source": source, **kwargs}))
+            except Exception:
+                self._pending_slots.release()
+                raise
+            return job
+
+    def cancel(self, job_id: str) -> bool:
+        """Request cancellation. True if the job was live.
+
+        Two cases, because they are genuinely different:
+
+        - **Not yet started** (queued, or the worker has not reached it): the job
+          is marked CANCELLED immediately and the worker skips it.
+        - **Running**: `cancel_requested` is set and the token is tripped. The
+          job stops at its next stage boundary and is then marked CANCELLED.
+          Until that boundary is reached it stays RUNNING with
+          `cancel_requested: true` - honest about work still in flight rather
+          than claiming an instant stop we cannot deliver.
+        """
+        job = self._store.get(job_id)
+        if job is None or job.is_terminal:
+            return False
+
+        with self._lock:
+            token = self._tokens.get(job_id)
+
+        if token is None:
+            self._store.update(
+                job_id,
+                state=JobState.CANCELLED,
+                progress="cancelled",
+                cancel_requested=True,
+            )
+        else:
+            token.cancel()
+            self._store.update(job_id, cancel_requested=True, progress="cancelling")
+        return True
+
+    def running_jobs(self) -> list[str]:
+        with self._lock:
+            return list(self._tokens)
+
+    @property
+    def pending_count(self) -> int:
+        return self._queue.qsize()
+
+    def _worker(self) -> None:
+        while True:
+            item = self._queue.get()
+            try:
+                if item is None:
+                    return
+                self._pending_slots.release()
+                job_id, kwargs = item
+
+                # Register the token BEFORE inspecting state, so a concurrent
+                # cancel() always finds a token to trip rather than racing us
+                # into marking the job cancelled while we start it anyway.
+                token = CancelToken()
+                with self._lock:
+                    self._tokens[job_id] = token
+                try:
+                    try:
+                        self._run_one(job_id, kwargs, token)
+                    except Exception as exc:  # noqa: BLE001 - keep the worker alive
+                        self._store.update(
+                            job_id,
+                            state=JobState.ERROR,
+                            error=f"{type(exc).__name__}: {exc}",
+                            progress="failed",
+                        )
+                finally:
+                    with self._lock:
+                        self._tokens.pop(job_id, None)
+            finally:
+                self._queue.task_done()
+
+    def _run_one(self, job_id: str, kwargs: dict[str, Any], token: CancelToken) -> None:
+        from textflowkit.core.runner import run_job
+
+        job = self._store.get(job_id)
+        if job is None or job.is_terminal:
+            return  # cancelled or reaped before we got to it
+        run_job(job, self._store, check_cancel=token.checkpoint, **kwargs)
+
+
+# Process-wide default executor, shared by the adapters.
+_default_executor: JobExecutor | None = None
+_executor_lock = threading.Lock()
+
+
+def get_default_executor() -> JobExecutor:
+    global _default_executor
+    with _executor_lock:
+        if _default_executor is None:
+            _default_executor = JobExecutor(get_default_store())
+        return _default_executor
+
+
+def reset_default_executor() -> None:
+    """Drop the cached executor (tests, embedding)."""
+    global _default_executor
+    with _executor_lock:
+        if _default_executor is not None:
+            _default_executor.shutdown(wait=False)
+        _default_executor = None