125 lines
4.9 KiB
TypeScript
125 lines
4.9 KiB
TypeScript
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, full-stopped 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);
|
|
});
|
|
});
|