193 lines
7.0 KiB
TypeScript
193 lines
7.0 KiB
TypeScript
import { describe, expect, test } from "bun:test";
|
||
import {
|
||
ParakeetVocabTokenizer,
|
||
parakeetFileName,
|
||
parakeetFileUrl,
|
||
parseVocab,
|
||
planAudioChunks,
|
||
planChunks,
|
||
transcribeWindows,
|
||
wordsToChunks,
|
||
type ParakeetWord,
|
||
} from "../src/lib/parakeet";
|
||
|
||
// Pure Parakeet engine logic (D-WS7a) — no model, no runtime.
|
||
|
||
const RATE = 16_000;
|
||
|
||
function word(text: string, start: number, end: number): ParakeetWord {
|
||
return { text, start_time: start, end_time: end };
|
||
}
|
||
|
||
function expectContiguous(windows: Array<[number, number]>, length: number) {
|
||
expect(windows[0]![0]).toBe(0);
|
||
expect(windows[windows.length - 1]![1]).toBe(length);
|
||
for (let i = 1; i < windows.length; i++) {
|
||
expect(windows[i]![0]).toBe(windows[i - 1]![1]);
|
||
}
|
||
}
|
||
|
||
describe("planChunks", () => {
|
||
const flat = () => 1;
|
||
|
||
test("audio up to the window length is a single window", () => {
|
||
expect(planChunks(10 * RATE, flat)).toEqual([[0, 10 * RATE]]);
|
||
expect(planChunks(30 * RATE, flat)).toEqual([[0, 30 * RATE]]);
|
||
expect(planChunks(0, flat)).toEqual([]);
|
||
});
|
||
|
||
test("windows never exceed 30 s, are contiguous, and cover everything", () => {
|
||
const length = 95 * RATE + 123;
|
||
const windows = planChunks(length, flat);
|
||
expectContiguous(windows, length);
|
||
for (const [start, end] of windows) {
|
||
expect(end - start).toBeLessThanOrEqual(30 * RATE);
|
||
expect(end).toBeGreaterThan(start);
|
||
}
|
||
});
|
||
|
||
test("a 90-minute file plans ≥ 180 bounded windows (no full-length copies needed)", () => {
|
||
const length = 90 * 60 * RATE;
|
||
const windows = planChunks(length, flat);
|
||
expect(windows.length).toBeGreaterThanOrEqual(180);
|
||
expectContiguous(windows, length);
|
||
for (const [start, end] of windows) expect(end - start).toBeLessThanOrEqual(30 * RATE);
|
||
});
|
||
|
||
test("cuts in the quietest 100 ms of the last 6 s before the limit", () => {
|
||
// Loud everywhere except a pause at 27.0–27.1 s.
|
||
const quietStart = 27 * RATE;
|
||
const energy = (start: number, end: number) =>
|
||
start >= quietStart && end <= quietStart + 0.1 * RATE ? 0 : 1;
|
||
const windows = planChunks(60 * RATE, energy);
|
||
const firstCut = windows[0]![1];
|
||
expect(firstCut).toBeGreaterThanOrEqual(quietStart);
|
||
expect(firstCut).toBeLessThanOrEqual(quietStart + 0.1 * RATE);
|
||
});
|
||
|
||
test("planAudioChunks finds a real silent gap in PCM", () => {
|
||
const audio = new Float32Array(50 * RATE);
|
||
for (let i = 0; i < audio.length; i++) audio[i] = Math.sin(i / 7) * 0.5;
|
||
audio.fill(0, 26 * RATE, 26.5 * RATE);
|
||
const windows = planAudioChunks(audio);
|
||
expect(windows[0]![1]).toBeGreaterThanOrEqual(26 * RATE);
|
||
expect(windows[0]![1]).toBeLessThanOrEqual(26.5 * RATE);
|
||
expectContiguous(windows, audio.length);
|
||
});
|
||
});
|
||
|
||
describe("wordsToChunks", () => {
|
||
test("splits after sentence-ending punctuation, keeping absolute timestamps", () => {
|
||
const chunks = wordsToChunks([
|
||
word("Hello", 0.5, 0.8),
|
||
word("there.", 0.9, 1.2),
|
||
word("How", 1.5, 1.7),
|
||
word("are", 1.7, 1.8),
|
||
word("you?", 1.8, 2.1),
|
||
]);
|
||
expect(chunks).toEqual([
|
||
{ text: " Hello there.", timestamp: [0.5, 1.2] },
|
||
{ text: " How are you?", timestamp: [1.5, 2.1] },
|
||
]);
|
||
});
|
||
|
||
test("closing quotes after the full stop still end the sentence", () => {
|
||
const chunks = wordsToChunks([word('"Yes."', 0, 0.4), word("Then", 0.5, 0.7)]);
|
||
expect(chunks.map((chunk) => chunk.text)).toEqual([' "Yes."', " Then"]);
|
||
});
|
||
|
||
test("a pause of at least 1 s starts a new chunk", () => {
|
||
const chunks = wordsToChunks([word("one", 0, 0.3), word("two", 1.4, 1.6)]);
|
||
expect(chunks).toHaveLength(2);
|
||
});
|
||
|
||
test("a run-on sentence is split once it reaches 12 s", () => {
|
||
const words = Array.from({ length: 40 }, (_, i) => word(`w${i}`, i * 0.5, i * 0.5 + 0.4));
|
||
const chunks = wordsToChunks(words);
|
||
expect(chunks.length).toBeGreaterThan(1);
|
||
for (const chunk of chunks) {
|
||
expect((chunk.timestamp[1] ?? 0) - chunk.timestamp[0]).toBeLessThanOrEqual(12.5);
|
||
}
|
||
expect(chunks.map((c) => c.text.trim()).join(" ").split(" ")).toHaveLength(40);
|
||
});
|
||
|
||
test("no words → no chunks", () => {
|
||
expect(wordsToChunks([])).toEqual([]);
|
||
});
|
||
});
|
||
|
||
describe("vocab tokenizer", () => {
|
||
const vocab = "<unk> 0\n▁the 1\n▁cat 2\ns 3\n. 4\n<blk> 5\n";
|
||
|
||
test("parseVocab indexes tokens by id", () => {
|
||
expect(parseVocab(vocab)).toEqual(["<unk>", "▁the", "▁cat", "s", ".", "<blk>"]);
|
||
});
|
||
|
||
test("decode joins pieces, skips blank, and attaches punctuation", () => {
|
||
const tokenizer = new ParakeetVocabTokenizer(parseVocab(vocab));
|
||
expect(tokenizer.blankId).toBe(5);
|
||
expect(tokenizer.blankToken).toBe("<blk>");
|
||
expect(tokenizer.decode([1, 5, 2, 3, 4])).toBe("the cats.");
|
||
});
|
||
});
|
||
|
||
describe("Parakeet files", () => {
|
||
test("int8 and fp16 file names; other dtypes are not offered", () => {
|
||
expect(parakeetFileName("encoder", "q8")).toBe("encoder-model.int8.onnx");
|
||
expect(parakeetFileName("decoder", "q8")).toBe("decoder_joint-model.int8.onnx");
|
||
expect(parakeetFileName("encoder", "fp16")).toBe("encoder-model.fp16.onnx");
|
||
expect(() => parakeetFileName("encoder", "q4")).toThrow();
|
||
});
|
||
|
||
test("URLs are pinned to the measured revision under the model id", () => {
|
||
expect(parakeetFileUrl("vocab.txt")).toBe(
|
||
"https://huggingface.co/ysdede/parakeet-tdt-0.6b-v3-onnx/resolve/" +
|
||
"e5356bf6e24bdb7a269536b0e9e7869bc034b893/vocab.txt",
|
||
);
|
||
});
|
||
});
|
||
|
||
describe("transcribeWindows", () => {
|
||
test("recognizes each window with its offset and reports progress + new chunks", async () => {
|
||
const audio = new Float32Array(70 * RATE).fill(0.1);
|
||
const offsets: number[] = [];
|
||
const sizes: number[] = [];
|
||
const progress: number[] = [];
|
||
const streamed: string[] = [];
|
||
const result = await transcribeWindows(
|
||
audio,
|
||
async (window, offset) => {
|
||
offsets.push(offset);
|
||
sizes.push(window.length);
|
||
return [word(`at${Math.round(offset)}.`, offset + 0.1, offset + 0.5)];
|
||
},
|
||
{
|
||
onWindow: (chunks, processed, total) => {
|
||
progress.push(processed);
|
||
expect(total).toBe(70);
|
||
streamed.push(...chunks.map((chunk) => chunk.text));
|
||
},
|
||
},
|
||
);
|
||
expect(offsets[0]).toBe(0);
|
||
expect(offsets.length).toBe(3);
|
||
for (const size of sizes) expect(size).toBeLessThanOrEqual(30 * RATE);
|
||
expect(progress[progress.length - 1]).toBe(70);
|
||
expect(progress).toEqual([...progress].sort((a, b) => a - b));
|
||
expect(result.chunks.map((c) => c.text)).toEqual(streamed);
|
||
expect(result.text).toBe(streamed.join("").trim());
|
||
expect(result.chunks[1]!.timestamp[0]).toBeCloseTo(offsets[1]! + 0.1, 5);
|
||
});
|
||
|
||
test("silence-only audio: empty transcript, progress still reaches the end", async () => {
|
||
const progress: number[] = [];
|
||
const result = await transcribeWindows(
|
||
new Float32Array(40 * RATE),
|
||
async () => [],
|
||
{ onWindow: (_chunks, processed) => progress.push(processed) },
|
||
);
|
||
expect(result).toEqual({ text: "", chunks: [] });
|
||
expect(progress[progress.length - 1]).toBe(40);
|
||
});
|
||
});
|