whisper_web/src/worker/longform/pipeline.ts

/**
 * Long-form file transcription (D-WS3a), mirroring faster-whisper's batched
 * pipeline without the batching: Silero VAD drops silence and cuts the audio
 * on speech boundaries into chunks ≤ 30 s; each chunk is one Whisper decode
 * (no overlapping windows to stitch); the language is detected once on the
 * first speech chunk and pinned; an optional quality check retries a bad
 * decode at higher temperatures.
 *
 * Memory (plan 001-0 Review Focus 1): VAD reads the caller's buffer in
 * place; only the chunk being decoded is copied (≤ 30 s), and the only state
 * that grows with the file is the transcript itself.
 */

import type { ProgressWindow, TranscriptChunk } from "../protocol";
import { TranscribeProgressTracker } from "../protocol";
import { compressionRatio, DEFAULT_FALLBACK, needsFallback, type FallbackOptions } from "./quality";
import type {
  DecodeRequest,
  DecodeResult,
  LongFormDeps,
  LongFormOptions,
  LongFormResult,
  RepetitionGuard,
} from "./types";
import {
  chunkAudio,
  chunkTimeToOriginal,
  collectChunks,
  getSpeechTimestamps,
  SAMPLE_RATE,
  type SpeechChunk,
} from "./vad";

/** Whisper's window: no chunk is longer. */
const MAX_CHUNK_SECONDS = 30;

/** No guard: Whisper's own decoding (the measured default, D-WS3a). */
export const DEFAULT_GUARD: RepetitionGuard = { noRepeatNgramSize: 0, repetitionPenalty: 1 };

/** The tuning the worker runs with — chosen by measurement (D-WS3a). */
export const LONG_FORM_DEFAULTS: Pick<LongFormOptions, "guard" | "fallback" | "vad"> = {
  guard: DEFAULT_GUARD,
  fallback: false,
};

export async function transcribeLongForm(
  audio: Float32Array,
  deps: LongFormDeps,
  options: LongFormOptions = {},
): Promise<LongFormResult> {
  const totalSeconds = audio.length / SAMPLE_RATE;
  options.onProgress?.(0, totalSeconds);

  const probs = await deps.speechProbabilities(audio);
  const spans = getSpeechTimestamps(probs, audio.length, {
    ...options.vad,
    maxSpeechDurationS: Math.min(options.vad?.maxSpeechDurationS ?? MAX_CHUNK_SECONDS, MAX_CHUNK_SECONDS),
  });
  const chunks = collectChunks(spans, MAX_CHUNK_SECONDS * SAMPLE_RATE);

  const tracker = TranscribeProgressTracker.forWindows(totalSeconds, chunks.map(progressWindow(chunks, totalSeconds)));
  const report = (processed: number | null) => {
    if (processed !== null) options.onProgress?.(processed, totalSeconds);
  };

  const guard = options.guard ?? DEFAULT_GUARD;
  const fallback: FallbackOptions | null =
    options.fallback === false || options.fallback === undefined
      ? null
      : { ...DEFAULT_FALLBACK, ...options.fallback };
  let language =
    options.language === undefined || options.language === "auto" ? undefined : options.language;

  const texts: string[] = [];
  const outChunks: TranscriptChunk[] = [];
  let decodeCalls = 0;
  let generatedTokens = 0;
  let speechSamples = 0;

  for (const chunk of chunks) {
    const chunkPcm = chunkAudio(audio, chunk);
    speechSamples += chunk.samples;
    if (language === undefined && deps.detectLanguage) {
      language = await deps.detectLanguage(chunkPcm);
    }
    const base: DecodeRequest = {
      ...(language !== undefined && { language }),
      ...(options.task !== undefined && { task: options.task }),
      ...(options.prompt !== undefined && { prompt: options.prompt }),
      temperature: 0,
      noRepeatNgramSize: guard.noRepeatNgramSize,
      repetitionPenalty: guard.repetitionPenalty,
      ...(options.onPartial && { onPartial: options.onPartial }),
      onTimestamp: (t) => report(tracker.timestamp(t)),
    };
    const attempts = fallback?.temperatures ?? [0];
    const tried: { result: DecodeResult; ratio: number }[] = [];
    let chosen: DecodeResult | null = null;
    for (const temperature of attempts) {
      const result = await deps.decode(chunkPcm, { ...base, temperature });
      decodeCalls += 1;
      generatedTokens += result.generatedTokens;
      if (!fallback) {
        chosen = result;
        break;
      }
      const ratio = await compressionRatio(result.text);
      if (!needsFallback({ compressionRatio: ratio, avgLogProb: result.avgLogProb }, fallback)) {
        chosen = result;
        break;
      }
      tried.push({ result, ratio });
    }
    chosen ??= bestAttempt(tried, fallback!);

    const text = chosen.text.trim();
    if (text.length > 0) texts.push(text);
    for (const segment of chosen.chunks) {
      const [start, end] = segment.timestamp;
      outChunks.push({
        text: segment.text,
        timestamp: [
          chunkTimeToOriginal(chunk, start),
          end === null ? null : chunkTimeToOriginal(chunk, end),
        ],
      });
    }
    report(tracker.finalizeWindow());
  }
  if (chunks.length === 0) report(tracker.finalizeWindow());

  return {
    text: texts.join(" "),
    chunks: outChunks,
    ...(language !== undefined && { language }),
    noSpeech: chunks.length === 0,
    speechSeconds: speechSamples / SAMPLE_RATE,
    decodeCalls,
    generatedTokens,
  };
}

/**
 * Every temperature failed: like faster-whisper, prefer attempts under the
 * compression threshold, then the highest average log-prob (unknown counts
 * as lowest), then the lowest compression ratio.
 */
function bestAttempt(
  tried: { result: DecodeResult; ratio: number }[],
  fallback: FallbackOptions,
): DecodeResult {
  const underThreshold = tried.filter((t) => t.ratio <= fallback.compressionRatioThreshold);
  const pool = underThreshold.length > 0 ? underThreshold : tried;
  return pool.reduce((best, t) => {
    const a = best.result.avgLogProb ?? -Infinity;
    const b = t.result.avgLogProb ?? -Infinity;
    if (b > a || (b === a && t.ratio < best.ratio)) return t;
    return best;
  }).result;
}

function progressWindow(chunks: SpeechChunk[], totalSeconds: number) {
  return (chunk: SpeechChunk, index: number): ProgressWindow => {
    const next = chunks[index + 1];
    const last = chunk.spans[chunk.spans.length - 1]!;
    return {
      toAbsolute: (t) => chunkTimeToOriginal(chunk, t),
      endSeconds: last.end / SAMPLE_RATE,
      nextStartSeconds: next ? next.spans[0]!.start / SAMPLE_RATE : totalSeconds,
    };
  };
}