Feed Whisper a decoder prompt through the ASR pipeline and strip it from the output
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Lucas Winther
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
942799d3cb
commit
e6de097cfd
@@ -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<string, unknown>): Promise<unknown>;
|
||||
_prepare_generation_config(config: null, kwargs: Record<string, unknown>): 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]);
|
||||
};
|
||||
}
|
||||
@@ -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<string, unknown>[] = [];
|
||||
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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user