231 lines
9.5 KiB
TypeScript
231 lines
9.5 KiB
TypeScript
import { describe, expect, test } from "bun:test";
|
||
import {
|
||
transcribeLongForm,
|
||
type DecodeRequest,
|
||
type DecodeResult,
|
||
type LongFormDeps,
|
||
} from "../src/worker/longform";
|
||
import { SAMPLE_RATE, VAD_WINDOW_SAMPLES } from "../src/worker/longform/vad";
|
||
|
||
const MAX_CHUNK_SAMPLES = 30 * SAMPLE_RATE;
|
||
|
||
/** Speech probabilities for audio of `seconds` with speech in [start, end) ranges. */
|
||
function maskProbs(seconds: number, speech: [number, number][]): Float32Array {
|
||
const windows = Math.ceil((seconds * SAMPLE_RATE) / VAD_WINDOW_SAMPLES);
|
||
const probs = new Float32Array(windows).fill(0.02);
|
||
for (const [start, end] of speech) {
|
||
const from = Math.floor((start * SAMPLE_RATE) / VAD_WINDOW_SAMPLES);
|
||
const to = Math.min(windows, Math.floor((end * SAMPLE_RATE) / VAD_WINDOW_SAMPLES));
|
||
for (let i = from; i < to; i++) probs[i] = 0.95;
|
||
}
|
||
return probs;
|
||
}
|
||
|
||
interface Recorder {
|
||
deps: LongFormDeps;
|
||
decodes: { audio: Float32Array; request: DecodeRequest }[];
|
||
vadInputs: Float32Array[];
|
||
detectCalls: Float32Array[];
|
||
}
|
||
|
||
function fakeDeps(
|
||
probs: Float32Array,
|
||
decode: (audio: Float32Array, request: DecodeRequest, call: number) => DecodeResult = (
|
||
audio,
|
||
) => ({
|
||
text: ` words ${audio.length}`,
|
||
chunks: [{ text: ` words ${audio.length}`, timestamp: [0, audio.length / SAMPLE_RATE] }],
|
||
generatedTokens: 5,
|
||
avgLogProb: -0.2,
|
||
}),
|
||
detect: string | null = null,
|
||
): Recorder {
|
||
const recorder: Recorder = { decodes: [], vadInputs: [], detectCalls: [], deps: null as never };
|
||
recorder.deps = {
|
||
speechProbabilities: async (audio) => {
|
||
recorder.vadInputs.push(audio);
|
||
return probs;
|
||
},
|
||
decode: async (audio, request) => {
|
||
recorder.decodes.push({ audio, request });
|
||
request.onPartial?.("x");
|
||
request.onTimestamp?.(1);
|
||
return decode(audio, request, recorder.decodes.length);
|
||
},
|
||
...(detect !== null && {
|
||
detectLanguage: async (audio: Float32Array) => {
|
||
recorder.detectCalls.push(audio);
|
||
return detect;
|
||
},
|
||
}),
|
||
};
|
||
return recorder;
|
||
}
|
||
|
||
describe("transcribeLongForm", () => {
|
||
test("no speech → empty result, decode never called", async () => {
|
||
const audio = new Float32Array(60 * SAMPLE_RATE);
|
||
const { deps, decodes } = fakeDeps(maskProbs(60, []));
|
||
const progress: number[] = [];
|
||
const result = await transcribeLongForm(audio, deps, {
|
||
onProgress: (processed) => progress.push(processed),
|
||
});
|
||
expect(decodes).toHaveLength(0);
|
||
expect(result.text).toBe("");
|
||
expect(result.chunks).toEqual([]);
|
||
expect(result.noSpeech).toBe(true);
|
||
expect(progress.at(-1)).toBe(60);
|
||
});
|
||
|
||
test("speech regions are decoded per chunk and timestamps mapped to the file", async () => {
|
||
// Two 5 s regions 15 s apart → one ≤ 30 s chunk holding both spans.
|
||
const audio = new Float32Array(40 * SAMPLE_RATE);
|
||
const { deps, decodes } = fakeDeps(maskProbs(40, [[5, 10], [25, 30]]), (chunk) => ({
|
||
text: " hello there",
|
||
chunks: [
|
||
{ text: " hello", timestamp: [0.5, 2] },
|
||
{ text: " there", timestamp: [6.5, 8] },
|
||
],
|
||
generatedTokens: 4,
|
||
avgLogProb: -0.1,
|
||
}));
|
||
const result = await transcribeLongForm(audio, deps, {});
|
||
expect(decodes).toHaveLength(1);
|
||
// Window-aligned spans padded by 400 ms: [73 472, 166 144) and
|
||
// [393 472, 486 144) → 2 × 92 672 samples, silence between them dropped.
|
||
expect(decodes[0]!.audio.length).toBe(2 * 92_672);
|
||
expect(result.text).toBe("hello there");
|
||
expect(result.noSpeech).toBe(false);
|
||
// The chunk starts at 4.592 s; 6.5 s into the chunk is past the first
|
||
// span (5.792 s long) → 0.708 s into the second, which starts at 24.592 s.
|
||
expect(result.chunks[0]!.timestamp[0]).toBeCloseTo(4.592 + 0.5, 3);
|
||
expect(result.chunks[1]!.timestamp[0]).toBeCloseTo(24.592 + 0.708, 3);
|
||
expect(result.chunks[1]!.timestamp[0]).toBeGreaterThan(24);
|
||
expect(result.chunks[1]!.timestamp[1]!).toBeLessThan(30.5);
|
||
});
|
||
|
||
test("auto language: detected once on the first speech chunk, then pinned", async () => {
|
||
const audio = new Float32Array(100 * SAMPLE_RATE);
|
||
const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 99]]), undefined, "da");
|
||
const result = await transcribeLongForm(audio, recorder.deps, { language: "auto" });
|
||
expect(recorder.detectCalls).toHaveLength(1);
|
||
expect(recorder.detectCalls[0]).toBe(recorder.decodes[0]!.audio);
|
||
expect(recorder.decodes.length).toBeGreaterThanOrEqual(3);
|
||
for (const { request } of recorder.decodes) expect(request.language).toBe("da");
|
||
expect(result.language).toBe("da");
|
||
});
|
||
|
||
test("an explicit language skips detection; English-only models send none", async () => {
|
||
const audio = new Float32Array(20 * SAMPLE_RATE);
|
||
const pinned = fakeDeps(maskProbs(20, [[2, 10]]), undefined, "de");
|
||
await transcribeLongForm(audio, pinned.deps, { language: "en" });
|
||
expect(pinned.detectCalls).toHaveLength(0);
|
||
expect(pinned.decodes[0]!.request.language).toBe("en");
|
||
|
||
const englishOnly = fakeDeps(maskProbs(20, [[2, 10]]));
|
||
await transcribeLongForm(audio, englishOnly.deps, {});
|
||
expect(englishOnly.decodes[0]!.request.language).toBeUndefined();
|
||
});
|
||
|
||
test("task, prompt and the repetition guard reach every decode", async () => {
|
||
const audio = new Float32Array(100 * SAMPLE_RATE);
|
||
const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 99]]));
|
||
await transcribeLongForm(audio, recorder.deps, {
|
||
task: "translate",
|
||
prompt: "Voices of the Void",
|
||
guard: { noRepeatNgramSize: 5, repetitionPenalty: 1.1 },
|
||
});
|
||
expect(recorder.decodes.length).toBeGreaterThanOrEqual(3);
|
||
for (const { request } of recorder.decodes) {
|
||
expect(request.task).toBe("translate");
|
||
expect(request.prompt).toBe("Voices of the Void");
|
||
expect(request.noRepeatNgramSize).toBe(5);
|
||
expect(request.repetitionPenalty).toBe(1.1);
|
||
expect(request.temperature).toBe(0);
|
||
}
|
||
});
|
||
|
||
test("fallback: a looping decode is retried at the next temperature", async () => {
|
||
const audio = new Float32Array(20 * SAMPLE_RATE);
|
||
const loop = " the main clar,".repeat(40);
|
||
const recorder = fakeDeps(maskProbs(20, [[2, 10]]), (_audio, request) => ({
|
||
text: request.temperature === 0 ? loop : " a fine sentence about festivals",
|
||
chunks: [],
|
||
generatedTokens: request.temperature === 0 ? 400 : 8,
|
||
avgLogProb: -0.2,
|
||
}));
|
||
const result = await transcribeLongForm(audio, recorder.deps, { fallback: {} });
|
||
expect(recorder.decodes.map((d) => d.request.temperature)).toEqual([0, 0.2]);
|
||
expect(result.text).toBe("a fine sentence about festivals");
|
||
expect(result.generatedTokens).toBe(408);
|
||
});
|
||
|
||
test("fallback: when every temperature fails, the best attempt is kept", async () => {
|
||
const audio = new Float32Array(20 * SAMPLE_RATE);
|
||
const recorder = fakeDeps(maskProbs(20, [[2, 10]]), (_audio, request) => ({
|
||
text: request.temperature === 0.6 ? " ok ok ok ok ok ok ok ok ok ok ok ok ok" : " la".repeat(80),
|
||
chunks: [],
|
||
generatedTokens: 10,
|
||
avgLogProb: null,
|
||
}));
|
||
const result = await transcribeLongForm(audio, recorder.deps, { fallback: {} });
|
||
expect(recorder.decodes).toHaveLength(6);
|
||
expect(result.text).toBe("ok ok ok ok ok ok ok ok ok ok ok ok ok");
|
||
});
|
||
|
||
test("fallback off → exactly one decode per chunk", async () => {
|
||
const audio = new Float32Array(20 * SAMPLE_RATE);
|
||
const recorder = fakeDeps(maskProbs(20, [[2, 10]]), () => ({
|
||
text: " la".repeat(80),
|
||
chunks: [],
|
||
generatedTokens: 10,
|
||
avgLogProb: -3,
|
||
}));
|
||
await transcribeLongForm(audio, recorder.deps, { fallback: false });
|
||
expect(recorder.decodes).toHaveLength(1);
|
||
});
|
||
|
||
test("progress is monotone and ends at the total; partials are forwarded", async () => {
|
||
const audio = new Float32Array(100 * SAMPLE_RATE);
|
||
const recorder = fakeDeps(maskProbs(100, [[0, 25], [40, 65], [75, 90]]));
|
||
const progress: number[] = [];
|
||
const partials: string[] = [];
|
||
await transcribeLongForm(audio, recorder.deps, {
|
||
onProgress: (processed, total) => {
|
||
expect(total).toBe(100);
|
||
progress.push(processed);
|
||
},
|
||
onPartial: (text) => partials.push(text),
|
||
});
|
||
expect(progress[0]).toBe(0);
|
||
expect(progress.at(-1)).toBe(100);
|
||
for (let i = 1; i < progress.length; i++) {
|
||
expect(progress[i]!).toBeGreaterThan(progress[i - 1]!);
|
||
}
|
||
expect(partials.length).toBe(recorder.decodes.length);
|
||
});
|
||
|
||
test("90-minute input: VAD sees the caller's buffer, decode sees ≤ 30 s chunks", async () => {
|
||
const seconds = 90 * 60;
|
||
const audio = new Float32Array(seconds * SAMPLE_RATE);
|
||
const speech: [number, number][] = [];
|
||
for (let t = 0; t < seconds; t += 25) speech.push([t, t + 20]);
|
||
const recorder = fakeDeps(maskProbs(seconds, speech));
|
||
const started = performance.now();
|
||
const result = await transcribeLongForm(audio, recorder.deps, {});
|
||
expect(performance.now() - started).toBeLessThan(5_000);
|
||
expect(recorder.vadInputs).toHaveLength(1);
|
||
expect(recorder.vadInputs[0]).toBe(audio); // no copy before VAD
|
||
let decoded = 0;
|
||
for (const { audio: chunk } of recorder.decodes) {
|
||
expect(chunk.length).toBeLessThanOrEqual(MAX_CHUNK_SAMPLES);
|
||
expect(chunk.buffer).not.toBe(audio.buffer);
|
||
decoded += chunk.length;
|
||
}
|
||
// Every speech second is decoded once (padding included), silence never.
|
||
expect(decoded).toBeLessThan(audio.length);
|
||
expect(result.speechSeconds).toBeCloseTo(decoded / SAMPLE_RATE, 3);
|
||
expect(result.decodeCalls).toBe(recorder.decodes.length);
|
||
});
|
||
});
|