From e6de097cfd3f19afe37cba9528e72087c0644519 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Thu, 29 Oct 2026 21:43:11 +0100 Subject: [PATCH] Feed Whisper a decoder prompt through the ASR pipeline and strip it from the output Co-Authored-By: Claude Opus 5.5 --- src/worker/prompt.ts | 101 +++++++++++++++++++++++++++++++++++ tests/prompt.test.ts | 124 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 225 insertions(+) create mode 100644 src/worker/prompt.ts create mode 100644 tests/prompt.test.ts diff --git a/src/worker/prompt.ts b/src/worker/prompt.ts new file mode 100644 index 0000000..c9514f4 --- /dev/null +++ b/src/worker/prompt.ts @@ -0,0 +1,101 @@ +/** + * Whisper decoder prompts (the "Names and terms" hint, D-WS2a) on top of the + * Transformers.js ASR pipeline. + * + * Transformers.js 4.2.0 (and 4.3.2) declare `prompt_ids` in the Whisper + * generation config but never use it — the parameter is commented out in + * `WhisperForConditionalGeneration.generate`. What generate() does honor is + * an explicit `decoder_input_ids`, so a prompt is expressed as + * `<|startofprev|> hint… <|startoftranscript|> <|lang|> <|task|>`. Two things + * then need fixing up around the call, which is what the wrapper does: + * + * - generate() returns the whole decoder sequence, and the pipeline's + * `_decode_asr` treats every non-special token as transcript text — the + * hint would be printed at the start of every window. The wrapper slices + * the prompt off the output again. + * - `max_new_tokens: 440` (D-021) assumes a 4-token prefix; with a prompt the + * total must still fit Whisper's 448 decoder positions. + * + * The prefix travels through the pipeline as a private kwarg: the ASR + * pipeline copies its call options into every per-window generate() call, so + * every 30 s window of a file and every live pass gets the same prompt, and + * concurrent passes cannot see each other's prompt (no shared mutable state). + * The `WhisperTextStreamer`'s `skip_prompt` drops the first `put()` — the + * whole decoder input, prompt included — so streamed partials stay clean. + */ + +import { fitHintToBudget } from "../lib/hint"; + +/** Private generate() kwarg carrying `[<|startofprev|>, ...hint ids]`. */ +export const PROMPT_PREFIX_KWARG = "whisper_web_prompt_prefix"; + +/** Whisper's decoder context (max_target_positions) for every catalog model. */ +export const WHISPER_DECODER_POSITIONS = 448; + +/** D-021's single-pass budget, used when a caller passes none. */ +const DEFAULT_MAX_NEW_TOKENS = 440; + +/** The slice of a Transformers.js Whisper model the wrapper relies on. */ +export interface PromptableModel { + generate(kwargs: Record): Promise; + _prepare_generation_config(config: null, kwargs: Record): unknown; + _retrieve_init_tokens(generationConfig: unknown): number[]; +} + +/** The slice of the Whisper tokenizer `buildPromptPrefix` needs. */ +export interface PromptTokenizer { + encode(text: string, options: { add_special_tokens: false }): number[]; + convert_tokens_to_ids(token: string): number; +} + +interface SliceableTensor { + slice(...slices: (number[] | null)[]): unknown; +} + +const WRAPPED = Symbol("whisper-web-prompt-support"); + +/** + * `[<|startofprev|>, ...ids]` for a raw user hint, truncated to Whisper's + * prompt budget (src/lib/hint.ts), or `[]` when the hint is empty. + */ +export function buildPromptPrefix(tokenizer: PromptTokenizer, hint: string): number[] { + const ids = fitHintToBudget(hint, (text) => + tokenizer.encode(text, { add_special_tokens: false }), + ); + if (ids.length === 0) return []; + return [tokenizer.convert_tokens_to_ids("<|startofprev|>"), ...ids]; +} + +/** + * Teach `model.generate` the `PROMPT_PREFIX_KWARG` option. Idempotent; calls + * without the kwarg (or with an empty prefix) are passed through untouched. + */ +export function installPromptSupport(model: PromptableModel): void { + const marked = model as PromptableModel & { [WRAPPED]?: true }; + if (marked[WRAPPED]) return; + marked[WRAPPED] = true; + + const original = model.generate.bind(model); + model.generate = async (kwargs) => { + const { [PROMPT_PREFIX_KWARG]: rawPrefix, ...rest } = kwargs; + const prefix = Array.isArray(rawPrefix) ? (rawPrefix as number[]) : []; + if (prefix.length === 0) return original(rest); + + // The same init tokens generate() would pick itself (language/task/ + // timestamps), with the prompt in front. + const initTokens = model._retrieve_init_tokens( + model._prepare_generation_config(null, rest), + ); + const decoderInputIds = [...prefix, ...initTokens]; + const requested = + typeof rest.max_new_tokens === "number" ? rest.max_new_tokens : DEFAULT_MAX_NEW_TOKENS; + const output = await original({ + ...rest, + decoder_input_ids: decoderInputIds, + max_new_tokens: Math.min(requested, WHISPER_DECODER_POSITIONS - decoderInputIds.length), + }); + // [batch=1, tokens] → drop the prompt columns; the sequence then starts + // at <|startoftranscript|> exactly as an unprompted pass would. + return (output as SliceableTensor).slice(null, [prefix.length, null]); + }; +} diff --git a/tests/prompt.test.ts b/tests/prompt.test.ts new file mode 100644 index 0000000..a91981e --- /dev/null +++ b/tests/prompt.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, test } from "bun:test"; +import { Tensor } from "@huggingface/transformers"; +import { + PROMPT_PREFIX_KWARG, + WHISPER_DECODER_POSITIONS, + buildPromptPrefix, + installPromptSupport, + type PromptableModel, +} from "../src/worker/prompt"; + +const PREV = 50362; + +/** + * Stand-in for a Transformers.js Whisper model: records the kwargs each + * generate() call received and, like the real one, returns the decoder input + * followed by the "generated" tokens as an int64 [1, n] tensor. + */ +function fakeModel(initTokens: number[], generated: number[] = [50365, 400, 401, 50400]) { + const calls: Record[] = []; + const model: PromptableModel & { calls: typeof calls } = { + calls, + _prepare_generation_config: (_config, kwargs) => ({ fromKwargs: kwargs }), + _retrieve_init_tokens: () => [...initTokens], + async generate(kwargs) { + calls.push(kwargs); + const input = (kwargs.decoder_input_ids as number[] | undefined) ?? initTokens; + const sequence = [...input, ...generated]; + return new Tensor("int64", BigInt64Array.from(sequence.map(BigInt)), [1, sequence.length]); + }, + }; + return model; +} + +function tokens(output: unknown): number[] { + return (output as Tensor).tolist()[0]!.map(Number); +} + +describe("installPromptSupport (D-WS2a)", () => { + test("without a prompt kwarg, generate is called exactly as before", async () => { + const model = fakeModel([50258, 50259, 50360]); + installPromptSupport(model); + const kwargs = { max_new_tokens: 440, return_timestamps: true }; + const output = await model.generate({ ...kwargs }); + expect(model.calls).toEqual([kwargs]); + expect(tokens(output)).toEqual([50258, 50259, 50360, 50365, 400, 401, 50400]); + }); + + test("an empty prefix is the same as no prompt", async () => { + const model = fakeModel([50258]); + installPromptSupport(model); + await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [] }); + expect(model.calls).toEqual([{ max_new_tokens: 440 }]); + }); + + test("a prefix is prepended to the init tokens and stripped from the output", async () => { + const model = fakeModel([50258, 50259, 50360]); + installPromptSupport(model); + const output = await model.generate({ + max_new_tokens: 440, + language: "en", + [PROMPT_PREFIX_KWARG]: [PREV, 1, 2], + }); + const call = model.calls[0]!; + expect(call.decoder_input_ids).toEqual([PREV, 1, 2, 50258, 50259, 50360]); + expect(call).not.toHaveProperty(PROMPT_PREFIX_KWARG); + expect(call.language).toBe("en"); + // Prompt tokens would otherwise reach _decode_asr as transcript text. + expect(tokens(output)).toEqual([50258, 50259, 50360, 50365, 400, 401, 50400]); + }); + + test("works for English-only models (init tokens are just <|startoftranscript|>)", async () => { + const model = fakeModel([50257]); + installPromptSupport(model); + const output = await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [50361, 7] }); + expect(model.calls[0]!.decoder_input_ids).toEqual([50361, 7, 50257]); + expect(tokens(output)[0]).toBe(50257); + }); + + test("caps max_new_tokens so prompt + init + new tokens fit Whisper's 448 positions", async () => { + const model = fakeModel([50258, 50259, 50360]); + installPromptSupport(model); + const prefix = [PREV, ...Array.from({ length: 223 }, (_, i) => 1000 + i)]; + await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: prefix }); + expect(WHISPER_DECODER_POSITIONS).toBe(448); + expect(model.calls[0]!.max_new_tokens).toBe(448 - 227); + // A short prompt keeps the D-021 value. + await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [PREV, 5] }); + expect(model.calls[1]!.max_new_tokens).toBe(440); + }); + + test("installing twice wraps once", async () => { + const model = fakeModel([50258]); + installPromptSupport(model); + installPromptSupport(model); + await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [PREV, 5] }); + expect(model.calls).toHaveLength(1); + expect(model.calls[0]!.decoder_input_ids).toEqual([PREV, 5, 50258]); + }); +}); + +describe("buildPromptPrefix", () => { + const tokenizer = { + encode: (text: string, _options: { add_special_tokens: false }) => + [...text].map((c) => c.charCodeAt(0)), + convert_tokens_to_ids: (token: string) => (token === "<|startofprev|>" ? PREV : -1), + }; + + test("no hint → no prefix", () => { + expect(buildPromptPrefix(tokenizer, "")).toEqual([]); + expect(buildPromptPrefix(tokenizer, " \n ")).toEqual([]); + }); + + test("a hint → <|startofprev|> then the space-led hint tokens", () => { + expect(buildPromptPrefix(tokenizer, "Max\nVoid")).toEqual([ + PREV, + ...[..." Max, Void"].map((c) => c.charCodeAt(0)), + ]); + }); + + test("a huge hint stays within the 223-token budget", () => { + const prefix = buildPromptPrefix(tokenizer, "Maxwell, ".repeat(500)); + expect(prefix.length).toBeLessThanOrEqual(224); + }); +});