whisper_web/src/worker/longform/transformers.ts

/**
 * `LongFormDeps` over the existing Transformers.js runtime (D-WS3a/D-WS3b):
 * Silero VAD runs as a plain ONNX model on the same onnxruntime instance the
 * Whisper pipeline uses, and Whisper decodes go through the already-loaded
 * ASR pipeline — no second inference runtime. Free of worker globals so
 * `bun test` drives it on onnxruntime-node.
 */

import {
  AutoModel,
  LogitsProcessor,
  LogitsProcessorList,
  Tensor,
  WhisperTextStreamer,
  type AutomaticSpeechRecognitionPipeline,
  type PretrainedModelOptions,
} from "@huggingface/transformers";
import { getDownloadedModels } from "../../lib/modelCache";
import type { TranscriptChunk } from "../protocol";
import type { DecodeRequest, DecodeResult, LongFormDeps } from "./types";
import { SAMPLE_RATE, VAD_WINDOW_SAMPLES } from "./vad";

/**
 * Silero VAD v5, ONNX export by onnx-community (MIT; `onnx/model.onnx`,
 * fp32, 2.2 MB). Fetched from the Hub like Whisper weights — CORS-enabled,
 * so COEP allows it — and kept in the same Cache API store (D-005, D-WS3b).
 */
export const VAD_MODEL_ID = "onnx-community/silero-vad";

/** Samples of the previous window prepended to each one (Silero v5 wrapper). */
const VAD_CONTEXT_SAMPLES = 64;

export interface SileroVad {
  speechProbabilities(audio: Float32Array): Promise<Float32Array>;
  dispose(): Promise<void>;
}

export interface LoadSileroVadOptions {
  /**
   * Never touch the network: load only if the model is already cached,
   * otherwise reject (offline, or never downloaded — Review Focus 3).
   */
  localOnly: boolean;
  /** onnxruntime device; the worker passes "wasm" (a 2 MB model gains nothing from a GPU). */
  device?: PretrainedModelOptions["device"];
  /** Is the model in the cache? Defaults to the browser Cache API probe (modelCache.ts). */
  isCached?: () => Promise<boolean>;
}

async function vadInBrowserCache(): Promise<boolean> {
  return (await getDownloadedModels([VAD_MODEL_ID])).has(VAD_MODEL_ID);
}

export async function loadSileroVad(options: LoadSileroVadOptions): Promise<SileroVad> {
  if (options.localOnly && !(await (options.isCached ?? vadInBrowserCache)())) {
    throw new Error("The speech-detection model is not downloaded (D-WS3b).");
  }
  const model = await AutoModel.from_pretrained(VAD_MODEL_ID, {
    // The repo ships no config.json: the model is a bare ONNX graph.
    config: { model_type: "custom" } as unknown as PretrainedModelOptions["config"],
    dtype: "fp32",
    ...(options.device !== undefined && { device: options.device }),
  });
  return {
    async speechProbabilities(audio) {
      const windows = Math.ceil(audio.length / VAD_WINDOW_SAMPLES);
      const probs = new Float32Array(windows);
      const sr = new Tensor("int64", [BigInt(SAMPLE_RATE)], []);
      let state = new Tensor("float32", new Float32Array(2 * 128), [2, 1, 128]);
      const input = new Float32Array(VAD_CONTEXT_SAMPLES + VAD_WINDOW_SAMPLES);
      for (let i = 0; i < windows; i++) {
        // [64 samples of context | 512-sample window], zero-padded at the end.
        input.copyWithin(0, VAD_WINDOW_SAMPLES);
        const window = audio.subarray(i * VAD_WINDOW_SAMPLES, (i + 1) * VAD_WINDOW_SAMPLES);
        input.set(window, VAD_CONTEXT_SAMPLES);
        input.fill(0, VAD_CONTEXT_SAMPLES + window.length);
        const out = (await model({
          input: new Tensor("float32", input, [1, input.length]),
          sr,
          state,
        })) as { output: Tensor; stateN: Tensor };
        probs[i] = (out.output.data as Float32Array)[0]!;
        state = out.stateN;
      }
      return probs;
    },
    dispose: async () => {
      await model.dispose();
    },
  };
}

// ---------------------------------------------------------------------------
// Whisper decodes
// ---------------------------------------------------------------------------

/**
 * Records the log-probability of each generated token and the token count.
 * Kept last in the processor list (after Whisper's timestamp rules), so it
 * sees the distribution the token is actually picked from. A step's token is
 * known on the next call (it is the last input id), so the final token (EOS
 * as a rule) is not averaged in — close enough for the −1.0 threshold.
 */
class LogProbRecorder extends LogitsProcessor {
  steps = 0;
  private sum = 0;
  private count = 0;
  private previous: { data: Float32Array; logSumExp: number } | null = null;

  override _call(inputIds: bigint[][], logits: Tensor): Tensor {
    const ids = inputIds[0]!;
    if (this.previous) {
      const token = Number(ids[ids.length - 1]);
      this.sum += this.previous.data[token]! - this.previous.logSumExp;
      this.count += 1;
    }
    const data = (logits as unknown as Tensor[])[0]!.data as Float32Array;
    let max = -Infinity;
    for (const v of data) if (v > max) max = v;
    let total = 0;
    for (const v of data) total += Math.exp(v - max);
    this.previous = { data, logSumExp: max + Math.log(total) };
    this.steps += 1;
    return logits;
  }

  average(): number | null {
    return this.count > 0 ? this.sum / this.count : null;
  }
}

/** A processor list that keeps `last` at the end whatever Whisper pushes. */
class ListWithLast extends LogitsProcessorList {
  constructor(last: LogitsProcessor) {
    super();
    this.processors.push(last);
  }

  override push(item: LogitsProcessor): void {
    this.processors.splice(this.processors.length - 1, 0, item);
  }
}

/** Only language tokens may be generated — one-step language detection. */
class LanguageOnly extends LogitsProcessor {
  constructor(private readonly allowed: Set<number>) {
    super();
  }

  override _call(_inputIds: bigint[][], logits: Tensor): Tensor {
    const data = (logits as unknown as Tensor[])[0]!.data as Float32Array;
    for (let i = 0; i < data.length; i++) if (!this.allowed.has(i)) data[i] = -Infinity;
    return logits;
  }
}

interface WhisperGenerationConfig {
  decoder_start_token_id: number;
  lang_to_id?: Record<string, number>;
}

export function createWhisperDeps(
  asr: AutomaticSpeechRecognitionPipeline,
  options: { multilingual: boolean },
): Pick<LongFormDeps, "decode" | "detectLanguage"> {
  const decode = async (audio: Float32Array, request: DecodeRequest): Promise<DecodeResult> => {
    const recorder = new LogProbRecorder();
    const generate: Record<string, unknown> = {
      return_timestamps: true,
      // One window, single-pass generate (D-021's seek-loop note applies).
      max_new_tokens: 440,
      logits_processor: new ListWithLast(recorder),
    };
    if (options.multilingual) {
      if (request.language !== undefined) generate.language = request.language;
      if (request.task !== undefined) generate.task = request.task;
    }
    if (request.noRepeatNgramSize > 0) generate.no_repeat_ngram_size = request.noRepeatNgramSize;
    if (request.repetitionPenalty !== 1) generate.repetition_penalty = request.repetitionPenalty;
    if (request.temperature > 0) {
      generate.do_sample = true;
      generate.temperature = request.temperature;
    }
    if (request.onPartial || request.onTimestamp) {
      generate.streamer = new WhisperTextStreamer(
        asr.tokenizer as ConstructorParameters<typeof WhisperTextStreamer>[0],
        {
          skip_prompt: true,
          callback_function: (text: string) => request.onPartial?.(text),
          on_chunk_start: (t: number) => request.onTimestamp?.(t),
          on_chunk_end: (t: number) => request.onTimestamp?.(t),
        },
      );
    }
    const output = await asr(audio, generate);
    const single = (Array.isArray(output) ? output[0] : output) as {
      text?: string;
      chunks?: TranscriptChunk[];
    };
    return {
      text: single.text ?? "",
      chunks: single.chunks ?? [],
      generatedTokens: recorder.steps,
      avgLogProb: recorder.average(),
    };
  };

  if (!options.multilingual) return { decode };

  const detectLanguage = async (audio: Float32Array): Promise<string> => {
    const model = asr.model as unknown as {
      generation_config: WhisperGenerationConfig;
      generate(args: Record<string, unknown>): Promise<Tensor>;
    };
    const config = model.generation_config;
    const languages = Object.entries(config.lang_to_id ?? {});
    const { input_features } = (await asr.processor(audio)) as { input_features: Tensor };
    const list = new LogitsProcessorList();
    list.push(new LanguageOnly(new Set(languages.map(([, id]) => id))));
    const sequences = await model.generate({
      inputs: input_features,
      decoder_input_ids: [config.decoder_start_token_id],
      max_new_tokens: 1,
      logits_processor: list,
    });
    const tokenId = Number((sequences.tolist() as bigint[][])[0]!.at(-1));
    const token = languages.find(([, id]) => id === tokenId)?.[0];
    if (!token) throw new Error(`Language detection produced a non-language token (${tokenId}).`);
    return token.slice(2, -2); // "<|da|>" → "da"
  };
  return { decode, detectLanguage };
}