@they-juanreina/compost-cli 0.1.1 → 0.1.2

package/transcriber/app/asr.py

@@ -0,0 +1,198 @@
+"""ASR wrapper (#10): Whisper-large-v3 with event-tag tokens via WhisperX.
+
+The heavy model (whisperx / faster-whisper / torch) is imported lazily so the
+service, the cue parser, and the test suite all work without the multi-GB
+weights installed. Install the `asr` extra and run inside the OrbStack
+container for real transcription:
+
+    pip install -e '.[asr]'
+
+`transcribe()` returns word-aligned utterances whose text may contain event
+tags; cue_parser.parse_transcript_cues() then lifts those into cues[].
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Any, Protocol
+
+
+@dataclass
+class ASRConfig:
+    model_name: str = "large-v3"
+    device: str = "auto"  # "cpu" | "cuda" | "mps" | "auto"
+    compute_type: str = "int8"
+    language: str | None = None
+    event_tags: bool = True
+    engine: str = "whisper"  # "whisper" (WhisperX, Docker/CPU) | "parakeet" (parakeet-mlx, native Metal)
+
+
+@dataclass
+class ASRResult:
+    utterances: list[dict[str, Any]] = field(default_factory=list)
+    language: str | None = None
+    model: str = ""
+
+
+class WhisperBackend(Protocol):
+    """Minimal surface the ASR wrapper needs. The real WhisperX backend
+    implements this; tests provide a fake."""
+
+    def transcribe(self, audio_path: str) -> dict[str, Any]: ...
+
+
+def build_whisperx_transcribe_kwargs(language: str | None) -> dict[str, Any]:
+    """Build the per-call kwargs for `whisperx.Model.transcribe()` so the
+    configured language hint reaches transcribe, not just load_model (#180).
+
+    Pre-fix, the hint was only passed to ``whisperx.load_model``; ``transcribe``
+    re-ran auto-detect per file, so the request-level ``"language":"en"`` was
+    effectively ignored. Tests the pure mapping without needing whisperx weights.
+    """
+    kwargs: dict[str, Any] = {"batch_size": 16}
+    if language:
+        kwargs["language"] = language
+    return kwargs
+
+
+class WhisperXBackend:  # pragma: no cover - needs multi-GB weights
+    """Concrete WhisperBackend wrapping `whisperx`.
+
+    Imports `whisperx` and `torch` lazily inside `__init__` so this module
+    remains importable in environments without the [asr] extra installed.
+    The constructor loads the model (multi-GB) the first time only — the
+    `_load_whisperx_backend` lru_cache ensures one instance per (model, device,
+    compute_type) tuple per process.
+    """
+
+    def __init__(self, config: ASRConfig):
+        try:
+            import torch  # type: ignore
+            import whisperx  # type: ignore
+        except ImportError as e:
+            raise RuntimeError(
+                "whisperx is not installed. Install the asr extra: pip install -e '.[asr]'"
+            ) from e
+
+        device = _resolve_device(config.device)
+        self._model = whisperx.load_model(
+            config.model_name,
+            device=device,
+            compute_type=config.compute_type,
+            language=config.language,
+            asr_options={"suppress_numerals": False},
+        )
+        self._align_model = None
+        self._align_metadata = None
+        self._device = device
+        self._whisperx = whisperx
+        self._torch = torch
+        # Pre-fix (#180): load_model() received the language hint but
+        # model.transcribe() didn't — WhisperX re-ran auto-detect per file
+        # ("No language specified, language will be detected ... (increases
+        # inference time)"). Hold the configured language on the backend and
+        # pass it through on every transcribe call so the hint actually skips
+        # the per-file detection step.
+        self._language = config.language
+
+    def transcribe(self, audio_path: str) -> dict[str, Any]:
+        audio = self._whisperx.load_audio(audio_path)
+        # Forward the configured language so WhisperX skips per-file auto-detect.
+        # When None, behavior is unchanged (auto-detect, then we use the result).
+        result = self._model.transcribe(audio, **build_whisperx_transcribe_kwargs(self._language))
+        language = result.get("language") or self._language or "en"
+
+        # Lazy-load the alignment model on first use (depends on detected language).
+        if self._align_model is None:
+            self._align_model, self._align_metadata = self._whisperx.load_align_model(
+                language_code=language, device=self._device
+            )
+
+        aligned = self._whisperx.align(
+            result["segments"],
+            self._align_model,
+            self._align_metadata,
+            audio,
+            self._device,
+            return_char_alignments=False,
+        )
+        return {"segments": aligned["segments"], "language": language}
+
+
+def _resolve_device(requested: str) -> str:  # pragma: no cover - env-dependent
+    """Map `auto` to the best available device. `cpu`/`cuda`/`mps` pass through."""
+    if requested != "auto":
+        return requested
+    try:
+        import torch  # type: ignore
+
+        if torch.cuda.is_available():
+            return "cuda"
+        if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+            return "mps"
+    except ImportError:
+        pass
+    return "cpu"
+
+
+@lru_cache(maxsize=1)
+def _load_whisperx_backend(config_key: str) -> WhisperBackend:  # pragma: no cover - needs weights
+    """Lazily construct the real WhisperX backend. Cached per process so the
+    multi-GB model loads once (cold-start cached)."""
+    # config_key encodes (model_name, device, compute_type); reconstruct.
+    model_name, device, compute_type = config_key.split(":", 2)
+    return WhisperXBackend(
+        ASRConfig(model_name=model_name, device=device, compute_type=compute_type)
+    )
+
+
+class Transcriber:
+    def __init__(self, config: ASRConfig | None = None, backend: WhisperBackend | None = None):
+        self.config = config or ASRConfig()
+        self._backend = backend
+
+    def _get_backend(self) -> WhisperBackend:
+        if self._backend is not None:
+            return self._backend
+        if self.config.engine == "parakeet":
+            from .asr_parakeet import _load_parakeet_backend, resolve_parakeet_model
+
+            return _load_parakeet_backend(
+                resolve_parakeet_model(self.config.model_name), self.config.language
+            )
+        key = f"{self.config.model_name}:{self.config.device}:{self.config.compute_type}"
+        return _load_whisperx_backend(key)
+
+    def transcribe(self, audio_path: str) -> ASRResult:
+        raw = self._get_backend().transcribe(audio_path)
+        utterances = _normalize_segments(raw.get("segments", []))
+        return ASRResult(
+            utterances=utterances,
+            language=raw.get("language", self.config.language),
+            model=self.config.model_name,
+        )
+
+
+def _normalize_segments(segments: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Convert backend segments into compost utterance dicts (pre-diarization,
+    pre-cue-extraction). Speaker ids are filled by the diarizer (#11)."""
+    out: list[dict[str, Any]] = []
+    for i, seg in enumerate(segments, start=1):
+        words = [
+            {"w": w["word"], "s": int(w["start"] * 1000), "e": int(w["end"] * 1000), "conf": w.get("score", 1.0)}
+            for w in seg.get("words", [])
+            if "start" in w and "end" in w
+        ]
+        out.append(
+            {
+                "id": f"U-{i:04d}",
+                "speaker_id": seg.get("speaker", "S?"),
+                "turn": i,
+                "start_ms": int(seg["start"] * 1000),
+                "end_ms": int(seg["end"] * 1000),
+                "text": seg.get("text", "").strip(),
+                "words": words,
+            }
+        )
+    return out

package/transcriber/app/asr_parakeet.py

@@ -0,0 +1,174 @@
+"""Parakeet (NVIDIA NeMo TDT) ASR backend via `parakeet-mlx` — the native
+Apple-Silicon (Metal) transcription path (#176).
+
+Conforms to the `WhisperBackend` Protocol in `asr.py`: `transcribe()` returns
+`{"segments": [...], "language": ...}` with per-word timestamps, so the rest of
+the pipeline (diarization, cue parsing, silence typing, prosody) is
+engine-agnostic and unchanged.
+
+`parakeet-mlx` requires Apple Silicon + MLX and is imported lazily, so this
+module stays importable (and the pure mapping helpers stay unit-testable)
+without the hardware or the ~2.5 GB weights. The heavy backend itself is marked
+`# pragma: no cover`, exactly like `WhisperXBackend`.
+
+Why Parakeet-TDT 0.6B v3 by default: it tops the Open ASR Leaderboard's
+convenient (local, Apple-Silicon, word-timestamped) tier — measured ~58.8x
+realtime on an M1 Max vs ~1.3x for WhisperX in the CPU container — with native
+frame-level word timestamps and 25-language (incl. Spanish) coverage.
+"""
+
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+from typing import Any
+
+from .asr import ASRConfig, WhisperBackend
+
+# A tiny EN-vs-ES language heuristic for the native Parakeet path (#190).
+# parakeet-mlx auto-detects internally but doesn't surface the detection in its
+# AlignedResult, so without a `--language` hint the transcript previously
+# recorded `language: "und"` (via `_detect_language`'s fallback). For a turnkey
+# zero-config run we'd rather record something sensible. Counts function-word
+# hits in the transcribed text; biased toward English when the signal is weak
+# (the v3 model is English-first multilingual).
+_ES_HEURISTIC_TOKENS = frozenset(
+    {
+        "que", "de", "la", "los", "las", "el", "en", "es", "no",
+        "una", "por", "con", "para", "del", "como", "pero", "más",
+    }
+)
+_EN_HEURISTIC_TOKENS = frozenset(
+    {
+        "the", "and", "of", "to", "in", "is", "for", "with",
+        "on", "that", "this", "you", "are", "was", "but", "they",
+    }
+)
+_WORD_RE = re.compile(r"[a-zA-Záéíóúñü]+")
+
+
+def guess_lang_from_text(text: str) -> str:
+    """Best-effort EN/ES guess for the Parakeet path when neither the model
+    nor a `--language` hint reveal the language. Falls back to ``en`` when the
+    text is empty or signal is too weak — never returns ``und`` (#190).
+    """
+    if not text:
+        return "en"
+    tokens = _WORD_RE.findall(text.lower())
+    if not tokens:
+        return "en"
+    es = sum(1 for t in tokens if t in _ES_HEURISTIC_TOKENS)
+    en = sum(1 for t in tokens if t in _EN_HEURISTIC_TOKENS)
+    # Require a clear ES margin to flip — otherwise default to EN (v3 is EN-first).
+    if es > max(en, 1) * 1.2:
+        return "es"
+    return "en"
+
+# Multilingual v3 (English + 24 European languages incl. Spanish) is the default;
+# v2 (`...-0.6b-v2`) is English-only with marginally better English WER.
+DEFAULT_PARAKEET_MODEL = "mlx-community/parakeet-tdt-0.6b-v3"
+
+# parakeet-mlx loads the whole file into a single Metal buffer unless chunked: a
+# 1-hour interview tries to allocate ~131 GB and blows past Metal's ~20 GB cap.
+# Chunk at 2 minutes (parakeet stitches chunks via its default 15 s overlap +
+# token timestamps) so arbitrarily long audio fits in memory. Tunable via env.
+DEFAULT_CHUNK_DURATION_S = 120.0
+
+
+def tokens_to_words(tokens: list[Any]) -> list[dict[str, Any]]:
+    """Merge parakeet sub-word tokens into words.
+
+    parakeet emits sub-word tokens where a word boundary is marked by a leading
+    space (e.g. ``[" If", ... ," we"," me","as"]`` → ``["If","we","measure"]``).
+    A token whose text starts with a space (or the first token) begins a new
+    word; the rest extend the current one. Timestamps are in **seconds** (the
+    pipeline's `_normalize_segments` converts to ms).
+    """
+    words: list[dict[str, Any]] = []
+    cur: dict[str, Any] | None = None
+    for t in tokens:
+        txt = getattr(t, "text", "")
+        if cur is None or txt.startswith(" "):
+            if cur is not None:
+                words.append(cur)
+            cur = {
+                "word": txt,
+                "start": float(t.start),
+                "end": float(t.end),
+                "score": float(getattr(t, "confidence", 1.0) or 1.0),
+            }
+        else:
+            cur["word"] += txt
+            cur["end"] = float(t.end)
+    if cur is not None:
+        words.append(cur)
+    for w in words:
+        w["word"] = w["word"].strip()
+    return [w for w in words if w["word"]]
+
+
+def result_to_segments(result: Any) -> list[dict[str, Any]]:
+    """Map a parakeet-mlx ``AlignedResult`` (sentences → tokens) to the
+    `WhisperBackend` segment shape: ``{start, end, text, words}`` (seconds)."""
+    segments: list[dict[str, Any]] = []
+    for sent in getattr(result, "sentences", None) or []:
+        segments.append(
+            {
+                "start": float(sent.start),
+                "end": float(sent.end),
+                "text": (getattr(sent, "text", "") or "").strip(),
+                "words": tokens_to_words(getattr(sent, "tokens", None) or []),
+            }
+        )
+    return segments
+
+
+def resolve_parakeet_model(model_name: str | None) -> str:
+    """A whisper-style model name (the ASRConfig default) means 'use the Parakeet
+    default'; an explicit parakeet id passes through."""
+    if model_name and "parakeet" in model_name:
+        return model_name
+    return DEFAULT_PARAKEET_MODEL
+
+
+class ParakeetMLXBackend:  # pragma: no cover - needs MLX + weights
+    """Concrete `WhisperBackend` wrapping `parakeet-mlx` (Apple Silicon / Metal)."""
+
+    def __init__(self, config: ASRConfig):
+        import os
+
+        try:
+            import parakeet_mlx  # type: ignore
+        except ImportError as e:
+            raise RuntimeError(
+                "parakeet-mlx is not installed (native Apple-Silicon ASR). "
+                "Install it in the native transcriber venv: pip install parakeet-mlx"
+            ) from e
+        self._model = parakeet_mlx.from_pretrained(resolve_parakeet_model(config.model_name))
+        self._language = config.language
+        self._chunk_s = float(os.environ.get("COMPOST_PARAKEET_CHUNK_S", DEFAULT_CHUNK_DURATION_S))
+
+    def transcribe(self, audio_path: str) -> dict[str, Any]:
+        # chunk_duration keeps long files within Metal's buffer cap (see above).
+        result = self._model.transcribe(audio_path, chunk_duration=self._chunk_s)
+        segments = result_to_segments(result)
+        # Language resolution priority (#190):
+        #   1. Whatever parakeet-mlx surfaces on the result (future-proof — the
+        #      API doesn't expose it today, but a future minor might).
+        #   2. The configured `--language` hint, if any.
+        #   3. A tiny EN/ES heuristic on the transcribed text — better than
+        #      letting `_detect_language` fall back to "und" on the zero-config
+        #      turnkey path.
+        language = getattr(result, "language", None) or self._language
+        if not language:
+            full_text = " ".join((s.get("text") or "") for s in segments)
+            language = guess_lang_from_text(full_text)
+        return {"segments": segments, "language": language}
+
+
+@lru_cache(maxsize=1)
+def _load_parakeet_backend(model_id: str, language: str | None) -> WhisperBackend:  # pragma: no cover
+    """Lazily construct + cache the Parakeet backend (one model load per process)."""
+    return ParakeetMLXBackend(
+        ASRConfig(model_name=model_id, language=language, engine="parakeet")
+    )

package/transcriber/app/cue_parser.py

@@ -0,0 +1,110 @@
+"""Cue parser (#10).
+
+Whisper-large-v3 with event-tag tokens emits inline markers like [laughter],
+[sigh], [cough], [clear_throat], [unintelligible], and code-switching markers.
+This module pulls those out of utterance text into structured cues[] entries
+(schema/cues.taxonomy.json) and returns the cleaned text.
+
+Pure and deterministic — no model. The ASR wrapper (asr.py) produces the
+tagged text; this turns it into cues.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any
+
+# Whisper/Whisper-AT event tag → compost cue kind (cues.taxonomy.json).
+TAG_TO_KIND: dict[str, str] = {
+    "laughter": "laughter",
+    "laugh": "laughter",
+    "laughs": "laughter",
+    "sigh": "sigh",
+    "sighs": "sigh",
+    "cough": "cough",
+    "coughs": "cough",
+    "clear_throat": "throat-clear",
+    "throat_clear": "throat-clear",
+    "throat-clear": "throat-clear",
+    "unintelligible": "unintelligible",
+    "inaudible": "unintelligible",
+    "code_switch": "code-switching",
+    "code-switch": "code-switching",
+    "code_switching": "code-switching",
+}
+
+# Default confidence assigned to a tag-derived cue when the ASR gives none.
+DEFAULT_CONFIDENCE = 0.8
+
+_TAG_RE = re.compile(r"\[([a-zA-Z_\-]+)\]")
+
+
+def _clean_text(text: str) -> str:
+    # Drop recognized event tags, collapse the resulting double spaces.
+    def repl(m: re.Match[str]) -> str:
+        return "" if m.group(1).lower() in TAG_TO_KIND else m.group(0)
+
+    return re.sub(r"\s{2,}", " ", _TAG_RE.sub(repl, text)).strip()
+
+
+def parse_cues_from_utterance(
+    utterance: dict[str, Any],
+    next_cue_index: int = 1,
+    confidence: float = DEFAULT_CONFIDENCE,
+) -> tuple[str, list[dict[str, Any]]]:
+    """Return (cleaned_text, cues) for one utterance.
+
+    Cue timing: if a word in `words[]` matches the tag, use that word's span;
+    otherwise fall back to the utterance span.
+    """
+    text = utterance.get("text", "")
+    words = utterance.get("words", [])
+    speaker_id = utterance.get("speaker_id")
+    cues: list[dict[str, Any]] = []
+    idx = next_cue_index
+
+    for m in _TAG_RE.finditer(text):
+        kind = TAG_TO_KIND.get(m.group(1).lower())
+        if kind is None:
+            continue
+        start_ms, end_ms = _tag_span(m.group(0), words, utterance)
+        cue: dict[str, Any] = {
+            "id": f"CUE-{idx:03d}",
+            "kind": kind,
+            "start_ms": start_ms,
+            "end_ms": end_ms,
+            "source": "audio",
+            "confidence": confidence,
+        }
+        if speaker_id is not None:
+            cue["speaker_id"] = speaker_id
+        cues.append(cue)
+        idx += 1
+
+    return _clean_text(text), cues
+
+
+def _tag_span(
+    tag_token: str,
+    words: list[dict[str, Any]],
+    utterance: dict[str, Any],
+) -> tuple[int, int]:
+    for w in words:
+        if w.get("w") == tag_token:
+            return int(w["s"]), int(w["e"])
+    return int(utterance["start_ms"]), int(utterance["end_ms"])
+
+
+def parse_transcript_cues(transcript: dict[str, Any]) -> dict[str, Any]:
+    """Extract cues from every utterance, append to cues[], strip tags from text.
+
+    Cue ids continue from any existing cues[]. Mutates and returns the transcript.
+    """
+    existing = transcript.setdefault("cues", [])
+    idx = len(existing) + 1
+    for utt in transcript.get("utterances", []):
+        cleaned, cues = parse_cues_from_utterance(utt, next_cue_index=idx)
+        utt["text"] = cleaned
+        existing.extend(cues)
+        idx += len(cues)
+    return transcript