feat(canary): split long audio at pauses to avoid the ~40s truncation

Canary-Qwen silently truncates audio past its ~40 s / 1024-token window, so
multi-minute recordings lose most of their transcript (validated: mean word
recovery vs Whisper falls from ~1.0 to ~0.28). Add a general pause-aware
segmentation utility (`segment_audios_at_pauses`, pure torch) that splits
over-window audio at silences so no word is cut, and wire Canary-Qwen to
chunk (default `greedy`), transcribe per chunk, and regroup transcripts in
time order into one ScriptLine per input. Short clips yield a single chunk,
so their behavior is unchanged.

Also forward the requested model revision to the worker so it loads the
pinned snapshot rather than defaulting to "main".

This is the Canary long-audio fix from #528, rebased directly onto alpha and
decoupled from the offline-loading work in #527 (which it was originally
stacked on).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: wilke0818

src/senselab/audio/tasks/preprocessing/__init__.py

@@ -1,3 +1,8 @@
 """.. include:: ./doc.md"""  # noqa: D415
 
 from .preprocessing import *  # noqa: F403
+from .silence_segmentation import (  # noqa: F401
+    SegmentStrategy,
+    pause_aware_boundaries,
+    segment_audios_at_pauses,
+)

src/senselab/audio/tasks/preprocessing/silence_segmentation.py

@@ -0,0 +1,258 @@
+"""Pause-aware segmentation: split audio into bounded chunks at silences.
+
+A general utility for any task that must feed a model audio no longer than some
+window (e.g. ASR backends with a fixed encoder/context limit such as Canary-Qwen
+or Granite): it splits audio that exceeds ``max_segment_s`` into ``<=
+max_segment_s`` pieces, choosing cut points **at pauses** (minimum short-time RMS
+energy) so no word is cut. Pure ``torch`` — no model or external dependency.
+
+Two strategies:
+
+* ``"greedy"`` — from the previous cut, take the *farthest* pause within
+  ``max_segment_s``. Minimizes the number of segments, but pins cuts near the cap
+  so cut quality is whatever pause happens to sit there.
+* ``"dp"`` — optimal segmentation: globally minimize
+  ``sum(cut_badness) + cut_penalty * n_cuts`` subject to every segment
+  ``<= max_segment_s``. Finds the deepest pauses anywhere and trades an extra
+  segment for a much cleaner cut when it is worth it (``cut_penalty`` is the dial).
+
+Both fall back to a forced minimum-energy cut across any pause-less run longer
+than ``max_segment_s``. ``"none"`` disables splitting (single full-length span).
+"""
+
+from __future__ import annotations
+
+from typing import List, Literal, Tuple
+
+import torch
+
+from senselab.audio.data_structures import Audio
+from senselab.audio.tasks.preprocessing.preprocessing import extract_segments
+
+SegmentStrategy = Literal["none", "greedy", "dp"]
+
+# Documented defaults (override per call):
+DEFAULT_MIN_PAUSE_S = 0.20  # minimum silence duration to qualify as a pause candidate
+DEFAULT_SILENCE_PERCENTILE = 15.0  # frames below this energy percentile are "quiet" (adaptive)
+DEFAULT_CUT_PENALTY = 0.2  # DP per-cut penalty in normalized-badness units [0, 1]
+DEFAULT_ENERGY_FRAME_S = 0.02  # short-time RMS frame size used to locate cut points
+
+# Tolerance for the ``<= max_segment_s`` feasibility comparisons. Forced-cut
+# positions are built by repeated addition of ``max_segment_s`` (see ``_dp``),
+# so a span that is exactly ``max_segment_s`` long can read as e.g.
+# ``0.30000000000000004`` and spuriously fail an exact ``> max_segment_s``
+# guard — severing the DP chain and collapsing the output to one oversized span.
+# 1e-6 s (well below one sample at any real rate, well above float-accumulation
+# error over thousands of frames) absorbs the rounding without moving any real
+# boundary.
+_FEASIBILITY_EPS = 1e-6
+
+
+def _rms_envelope(audio: Audio, frame_s: float) -> Tuple[torch.Tensor, int, int, int]:
+    """Return (per-frame RMS, hop samples, sampling_rate, n_samples)."""
+    sr = int(audio.sampling_rate)
+    mono = audio.waveform.mean(dim=0)
+    n_samples = int(mono.shape[0])
+    hop = max(1, int(round(frame_s * sr)))
+    n_frames = n_samples // hop
+    if n_frames < 1:
+        return torch.zeros(0), hop, sr, n_samples
+    rms = mono[: n_frames * hop].reshape(n_frames, hop).pow(2).mean(dim=1).sqrt()
+    return rms, hop, sr, n_samples
+
+
+def _normalized(rms: torch.Tensor) -> torch.Tensor:
+    """Map RMS to [0, 1] badness via the 95th percentile (speech~1, silence~0)."""
+    if rms.numel() == 0:
+        return rms
+    ref = float(torch.quantile(rms, 0.95).item()) or float(rms.max().item()) or 1.0
+    return (rms / ref).clamp(0.0, 1.0)
+
+
+def _frame_time(frame_idx: int, hop: int, sr: int) -> float:
+    """Center time (seconds) of a frame."""
+    return (frame_idx * hop + hop / 2) / sr
+
+
+def _pause_candidates(
+    rms: torch.Tensor, hop: int, sr: int, min_pause_s: float, silence_percentile: float
+) -> List[Tuple[float, float]]:
+    """Return ``(time_s, badness)`` at the deepest frame of each detected pause."""
+    if rms.numel() == 0:
+        return []
+    norm = _normalized(rms)
+    thr = float(torch.quantile(rms, silence_percentile / 100.0).item())
+    # Materialize the mask to a Python list once: element-by-element ``bool(quiet[i])``
+    # on a tensor is very slow over the many frames of a long recording.
+    quiet = (rms <= thr).tolist()
+    min_frames = max(1, int(round(min_pause_s * sr / hop)))
+    out: List[Tuple[float, float]] = []
+    n = len(quiet)
+    i = 0
+    while i < n:
+        if not quiet[i]:
+            i += 1
+            continue
+        j = i
+        while j < n and quiet[j]:
+            j += 1
+        if (j - i) >= min_frames:
+            k = i + int(rms[i:j].argmin().item())
+            out.append((_frame_time(k, hop, sr), float(norm[k].item())))
+        i = j
+    return out
+
+
+def _forced_cut(rms: torch.Tensor, hop: int, sr: int, lo_s: float, hi_s: float) -> Tuple[float, float]:
+    """Quietest ``(time_s, badness)`` in ``[lo_s, hi_s]`` — fallback when no pause exists."""
+    norm = _normalized(rms)
+    lo_f = max(0, int(lo_s * sr / hop))
+    hi_f = min(rms.shape[0], max(lo_f + 1, int(hi_s * sr / hop)))
+    seg = rms[lo_f:hi_f]
+    if seg.numel() == 0:
+        return hi_s, 1.0
+    k = lo_f + int(seg.argmin().item())
+    return _frame_time(k, hop, sr), float(norm[k].item())
+
+
+def _spans_from_cuts(cuts: List[float], duration: float) -> List[Tuple[float, float]]:
+    pts = sorted({0.0, *[c for c in cuts if 0.0 < c < duration], duration})
+    return [(pts[k], pts[k + 1]) for k in range(len(pts) - 1)]
+
+
+def _greedy(
+    rms: torch.Tensor, hop: int, sr: int, duration: float, max_seg: float, cand_times: List[float]
+) -> List[float]:
+    cuts: List[float] = []
+    prev = 0.0
+    while duration - prev > max_seg:
+        feasible = [t for t in cand_times if prev < t <= prev + max_seg]
+        if feasible:
+            cut = max(feasible)
+        else:
+            cut, _ = _forced_cut(rms, hop, sr, prev + 0.5 * max_seg, prev + max_seg)
+            cut = min(max(cut, prev + 1.0), prev + max_seg)
+        cuts.append(cut)
+        prev = cut
+    return cuts
+
+
+def _dp(
+    rms: torch.Tensor,
+    hop: int,
+    sr: int,
+    duration: float,
+    max_seg: float,
+    candidates: List[Tuple[float, float]],
+    cut_penalty: float,
+) -> List[float]:
+    # Candidate positions: start, pauses, end. Insert forced cuts so no two
+    # consecutive candidates are farther than max_seg apart (guarantees feasibility).
+    raw = [(0.0, 0.0)] + [(t, b) for t, b in candidates if 0.0 < t < duration] + [(duration, 0.0)]
+    raw.sort()
+    pts: List[Tuple[float, float]] = [raw[0]]
+    for idx in range(1, len(raw)):
+        while raw[idx][0] - pts[-1][0] > max_seg + _FEASIBILITY_EPS:
+            ft, fb = _forced_cut(rms, hop, sr, pts[-1][0] + 0.5 * max_seg, pts[-1][0] + max_seg)
+            ft = min(max(ft, pts[-1][0] + 1.0), pts[-1][0] + max_seg)
+            pts.append((ft, fb))
+        pts.append(raw[idx])
+
+    m = len(pts)
+    inf = float("inf")
+    cost = [inf] * m
+    back = [-1] * m
+    cost[0] = 0.0
+    for i in range(1, m):
+        ti, bi = pts[i]
+        is_end = i == m - 1
+        for j in range(i - 1, -1, -1):
+            if ti - pts[j][0] > max_seg + _FEASIBILITY_EPS:
+                break
+            add = 0.0 if is_end else (bi + cut_penalty)  # endpoint is not a cut
+            c = cost[j] + add
+            if c < cost[i]:
+                cost[i] = c
+                back[i] = j
+
+    cuts: List[float] = []
+    i = m - 1
+    while i > 0:
+        cuts.append(pts[i][0])
+        i = back[i]
+    return cuts
+
+
+def pause_aware_boundaries(
+    audio: Audio,
+    max_segment_s: float,
+    strategy: SegmentStrategy = "dp",
+    *,
+    min_pause_s: float = DEFAULT_MIN_PAUSE_S,
+    silence_percentile: float = DEFAULT_SILENCE_PERCENTILE,
+    cut_penalty: float = DEFAULT_CUT_PENALTY,
+    energy_frame_s: float = DEFAULT_ENERGY_FRAME_S,
+) -> List[Tuple[float, float]]:
+    """Return ``(start, end)`` second-spans tiling ``audio`` with each span <= ``max_segment_s``.
+
+    Cuts are placed at pauses (see module docstring). Audio at or under
+    ``max_segment_s`` (or ``strategy="none"``) returns a single full-length span.
+    """
+    if max_segment_s <= 0:
+        raise ValueError(f"max_segment_s must be strictly positive, got {max_segment_s!r}.")
+    if min_pause_s <= 0:
+        raise ValueError(f"min_pause_s must be strictly positive, got {min_pause_s!r}.")
+    if energy_frame_s <= 0:
+        raise ValueError(f"energy_frame_s must be strictly positive, got {energy_frame_s!r}.")
+    if not 0.0 <= silence_percentile <= 100.0:
+        raise ValueError(f"silence_percentile must be in [0, 100], got {silence_percentile!r}.")
+    if cut_penalty < 0:
+        raise ValueError(f"cut_penalty must be non-negative, got {cut_penalty!r}.")
+
+    duration = audio.waveform.shape[1] / audio.sampling_rate
+    if strategy == "none":
+        return [(0.0, duration)]
+    if strategy not in ("greedy", "dp"):
+        raise ValueError(f"Unknown segmentation strategy {strategy!r}; expected 'none', 'greedy', or 'dp'.")
+
+    rms, hop, sr, n = _rms_envelope(audio, energy_frame_s)
+    if duration <= max_segment_s or rms.numel() == 0:
+        return [(0.0, duration)]
+    candidates = _pause_candidates(rms, hop, sr, min_pause_s, silence_percentile)
+
+    if strategy == "greedy":
+        cuts = _greedy(rms, hop, sr, duration, max_segment_s, [t for t, _ in candidates])
+    else:
+        cuts = _dp(rms, hop, sr, duration, max_segment_s, candidates, cut_penalty)
+    return _spans_from_cuts(cuts, duration)
+
+
+def segment_audios_at_pauses(
+    audios: List[Audio],
+    max_segment_s: float,
+    strategy: SegmentStrategy = "dp",
+    *,
+    min_pause_s: float = DEFAULT_MIN_PAUSE_S,
+    silence_percentile: float = DEFAULT_SILENCE_PERCENTILE,
+    cut_penalty: float = DEFAULT_CUT_PENALTY,
+    energy_frame_s: float = DEFAULT_ENERGY_FRAME_S,
+) -> List[List[Audio]]:
+    """Split each input audio at pauses into ``<= max_segment_s`` sub-Audios.
+
+    Returns, per input audio, the list of its sub-Audios in time order. An audio
+    that needs no splitting is returned as a single-element list holding the
+    original object (no copy).
+    """
+    out: List[List[Audio]] = []
+    for audio in audios:
+        spans = pause_aware_boundaries(
+            audio,
+            max_segment_s,
+            strategy,
+            min_pause_s=min_pause_s,
+            silence_percentile=silence_percentile,
+            cut_penalty=cut_penalty,
+            energy_frame_s=energy_frame_s,
+        )
+        out.append([audio] if len(spans) == 1 else extract_segments([(audio, spans)])[0])
+    return out