whisper_web/tests/dsp.ts

/**
 * Shared DSP helpers for the live-accuracy and resampler tests.
 *
 * The accuracy harness encodes "words" as pure tones at word-specific
 * frequencies and decodes them with a Goertzel filter bank ("oracle engine").
 * That makes transcript accuracy of the live *pipeline* measurable under
 * `bun test` with no model: if the segment loop cuts a word in half, drops
 * audio, duplicates a boundary, leaks stale text, or the resampler shifts /
 * aliases frequencies, the decoded word sequence changes and WER rises above
 * zero.
 */

import type { LiveTranscribeResult } from "../src/lib/live";

export const RATE = 16_000;

// ---------------------------------------------------------------------------
// Signal generation
// ---------------------------------------------------------------------------

export function sine(
  freq: number,
  seconds: number,
  rate: number = RATE,
  amplitude = 0.2,
): Float32Array {
  const out = new Float32Array(Math.round(seconds * rate));
  for (let i = 0; i < out.length; i++) {
    out[i] = amplitude * Math.sin((2 * Math.PI * freq * i) / rate);
  }
  return out;
}

export function silence(seconds: number, rate: number = RATE): Float32Array {
  return new Float32Array(Math.round(seconds * rate));
}

export function concatAudio(parts: readonly Float32Array[]): Float32Array {
  const out = new Float32Array(parts.reduce((n, p) => n + p.length, 0));
  let offset = 0;
  for (const part of parts) {
    out.set(part, offset);
    offset += part.length;
  }
  return out;
}

// ---------------------------------------------------------------------------
// Tone-word vocabulary and session builder
// ---------------------------------------------------------------------------

export const ORACLE_WORDS = [
  "alpha",
  "bravo",
  "charlie",
  "delta",
  "echo",
  "foxtrot",
  "golf",
  "hotel",
] as const;

/** Word k ↔ tone at 500 + 250·k Hz (500–2250 Hz, well inside speech band). */
export function wordFreq(index: number): number {
  return 500 + 250 * index;
}

export const WORD_SECONDS = 0.25;
export const GAP_SECONDS = 0.08;

export interface BuiltSession {
  audio: Float32Array;
  /** Ground truth: every word with its absolute onset time in seconds. */
  truth: { word: string; start: number }[];
  /** The reference transcript (all words, in order, space-joined). */
  reference: string;
}

/**
 * Lay out utterances (word lists) separated by `pauseSeconds` of silence,
 * with `GAP_SECONDS` between words inside an utterance.
 */
export function buildSession(
  utterances: readonly (readonly string[])[],
  options: {
    rate?: number;
    leadSeconds?: number;
    pauseSeconds?: number;
    /**
     * Duration of the very first word. Segments start relative to speech
     * onset, so stretching the first word is how a test shifts where a
     * max-length cap lands within the word/gap cycle (a longer lead-in
     * silence shifts onset and cap together and changes nothing).
     */
    firstWordSeconds?: number;
  } = {},
): BuiltSession {
  const rate = options.rate ?? RATE;
  const lead = options.leadSeconds ?? 0.5;
  const pause = options.pauseSeconds ?? 2.0;

  const parts: Float32Array[] = [silence(lead, rate)];
  const truth: { word: string; start: number }[] = [];
  let t = lead;
  let first = true;
  utterances.forEach((words, u) => {
    words.forEach((word, w) => {
      const index = ORACLE_WORDS.indexOf(word as (typeof ORACLE_WORDS)[number]);
      if (index === -1) throw new Error(`unknown oracle word: ${word}`);
      const duration = first ? (options.firstWordSeconds ?? WORD_SECONDS) : WORD_SECONDS;
      first = false;
      truth.push({ word, start: t });
      parts.push(sine(wordFreq(index), duration, rate));
      t += duration;
      if (w < words.length - 1) {
        parts.push(silence(GAP_SECONDS, rate));
        t += GAP_SECONDS;
      }
    });
    if (u < utterances.length - 1) {
      parts.push(silence(pause, rate));
      t += pause;
    }
  });
  return {
    audio: concatAudio(parts),
    truth,
    reference: truth.map((entry) => entry.word).join(" "),
  };
}

// ---------------------------------------------------------------------------
// Goertzel filter bank + oracle decoder
// ---------------------------------------------------------------------------

/** Normalized power of `freq` over samples[from, to). */
export function goertzelPower(
  samples: Float32Array,
  from: number,
  to: number,
  freq: number,
  rate: number,
): number {
  const coeff = 2 * Math.cos((2 * Math.PI * freq) / rate);
  let s1 = 0;
  let s2 = 0;
  for (let i = from; i < to; i++) {
    const s0 = samples[i]! + coeff * s1 - s2;
    s2 = s1;
    s1 = s0;
  }
  const n = to - from;
  return (s1 * s1 + s2 * s2 - coeff * s1 * s2) / (n * n);
}

const ORACLE_FRAME_SECONDS = 0.02;
/** Runs shorter than this are boundary noise and ignored. */
const ORACLE_FRAGMENT_MIN_SECONDS = 0.06;
/** Runs at least this long are a complete word (words are 0.25 s). */
const ORACLE_WORD_MIN_SECONDS = 0.22;
/** What a truncated word decodes as — a substitution, like real Whisper. */
export const ORACLE_FRAGMENT = "▯";

/**
 * Decode tone-words from 16 kHz audio. Deterministic stand-in for Whisper:
 * frames are labeled by dominant vocabulary frequency; runs ≥ 220 ms decode
 * as the word, and truncated runs (60–220 ms — a word cut by a segment
 * boundary) decode as the garbage token `ORACLE_FRAGMENT`, because a real
 * model mis-transcribes a cut word rather than silently skipping it. Either
 * fragment of a mid-word cut therefore raises WER.
 * Output mimics Whisper: leading-space text, per-word timestamped chunks.
 */
export function oracleDecode(audio: Float32Array): LiveTranscribeResult {
  const hop = Math.round(RATE * ORACLE_FRAME_SECONDS);
  const fragmentRun = Math.round(ORACLE_FRAGMENT_MIN_SECONDS / ORACLE_FRAME_SECONDS);
  const wordRun = Math.round(ORACLE_WORD_MIN_SECONDS / ORACLE_FRAME_SECONDS);

  const labels: number[] = [];
  for (let start = 0; start + hop <= audio.length; start += hop) {
    let energy = 0;
    for (let i = start; i < start + hop; i++) energy += audio[i]! * audio[i]!;
    if (Math.sqrt(energy / hop) < 0.05) {
      labels.push(-1);
      continue;
    }
    let best = -1;
    let bestPower = 0;
    for (let k = 0; k < ORACLE_WORDS.length; k++) {
      const power = goertzelPower(audio, start, start + hop, wordFreq(k), RATE);
      if (power > bestPower) {
        bestPower = power;
        best = k;
      }
    }
    labels.push(best);
  }

  const words: string[] = [];
  const chunks: LiveTranscribeResult["chunks"] = [];
  let runLabel = -1;
  let runStart = 0;
  const endRun = (endFrame: number) => {
    const length = endFrame - runStart;
    if (runLabel < 0 || length < fragmentRun) return;
    const word = length >= wordRun ? ORACLE_WORDS[runLabel]! : ORACLE_FRAGMENT;
    words.push(word);
    chunks.push({
      text: ` ${word}`,
      timestamp: [runStart * ORACLE_FRAME_SECONDS, endFrame * ORACLE_FRAME_SECONDS],
    });
  };
  labels.forEach((label, frame) => {
    if (label !== runLabel) {
      endRun(frame);
      runLabel = label;
      runStart = frame;
    }
  });
  endRun(labels.length);

  return { text: words.length > 0 ? ` ${words.join(" ")}` : "", chunks };
}

// ---------------------------------------------------------------------------
// Word error rate
// ---------------------------------------------------------------------------

/**
 * Standard WER: word-level Levenshtein distance over the reference length.
 * 0 = perfect; insertions (duplicated boundary fragments), deletions
 * (dropped audio), and substitutions all count.
 */
export function wordErrorRate(reference: string, hypothesis: string): number {
  const ref = reference.trim().split(/\s+/).filter(Boolean);
  const hyp = hypothesis.trim().split(/\s+/).filter(Boolean);
  let previous = Array.from({ length: hyp.length + 1 }, (_, j) => j);
  for (let i = 1; i <= ref.length; i++) {
    const current = new Array<number>(hyp.length + 1);
    current[0] = i;
    for (let j = 1; j <= hyp.length; j++) {
      current[j] = Math.min(
        previous[j]! + 1,
        current[j - 1]! + 1,
        previous[j - 1]! + (ref[i - 1] === hyp[j - 1] ? 0 : 1),
      );
    }
    previous = current;
  }
  return previous[hyp.length]! / Math.max(1, ref.length);
}