Add Parakeet window planning, word-to-chunk mapping, vocab tokenizer and pinned asset URLs
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
d554162a2c
commit
c9fdab1726
@@ -0,0 +1,230 @@
|
||||
/**
|
||||
* Pure logic of the Parakeet TDT engine path (Decisions.md D-WS7a): window
|
||||
* planning, word → transcript-chunk mapping, the vocab tokenizer, and the
|
||||
* pinned asset URLs. No runtime, no browser globals — the worker side
|
||||
* (src/worker/parakeetEngine.ts) injects the actual recognizer.
|
||||
*
|
||||
* Why our own windowing instead of parakeet.js's `transcribeLongAudio`: that
|
||||
* helper clamps windows to 20–180 s (90 s by default), while the spec caps
|
||||
* one encoder call at ~35 s (encoder attention memory grows with length).
|
||||
* Windows here are ≤ 30 s and end in the quietest 100 ms near the limit, so
|
||||
* words are almost never cut and no overlap merging is needed.
|
||||
*/
|
||||
|
||||
import type { TranscriptChunk, WhisperDtype } from "../worker/protocol";
|
||||
import { PARAKEET_MODEL_ID, PARAKEET_REVISION } from "./models";
|
||||
|
||||
/** Sample rate of all app audio (D-012) and of Parakeet's feature extractor. */
|
||||
export const PARAKEET_SAMPLE_RATE = 16_000;
|
||||
|
||||
/** One recognized word with absolute times in seconds (parakeet.js shape). */
|
||||
export interface ParakeetWord {
|
||||
text: string;
|
||||
start_time: number;
|
||||
end_time: number;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Window planning
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface PlanChunksOptions {
|
||||
sampleRate?: number;
|
||||
/** Hard cap on one window (one encoder call). Default 30 s. */
|
||||
maxSeconds?: number;
|
||||
/** How far back from the cap to look for a quiet cut point. Default 6 s. */
|
||||
searchSeconds?: number;
|
||||
/** Energy frame size for the cut search. Default 0.1 s. */
|
||||
frameSeconds?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Split `length` samples into contiguous `[start, end)` windows of at most
|
||||
* `maxSeconds`. Each window that isn't the last ends at the centre of the
|
||||
* lowest-energy frame in its final `searchSeconds` (ties → the later frame,
|
||||
* keeping windows long). `energyAt(start, end)` scores a sample range.
|
||||
*/
|
||||
export function planChunks(
|
||||
length: number,
|
||||
energyAt: (start: number, end: number) => number,
|
||||
options: PlanChunksOptions = {},
|
||||
): Array<[number, number]> {
|
||||
const rate = options.sampleRate ?? PARAKEET_SAMPLE_RATE;
|
||||
const maxSamples = Math.round((options.maxSeconds ?? 30) * rate);
|
||||
const searchSamples = Math.round((options.searchSeconds ?? 6) * rate);
|
||||
const frame = Math.max(1, Math.round((options.frameSeconds ?? 0.1) * rate));
|
||||
|
||||
const windows: Array<[number, number]> = [];
|
||||
let start = 0;
|
||||
while (start < length) {
|
||||
if (length - start <= maxSamples) {
|
||||
windows.push([start, length]);
|
||||
break;
|
||||
}
|
||||
const limit = start + maxSamples;
|
||||
let cut = limit;
|
||||
let best = Number.POSITIVE_INFINITY;
|
||||
for (let f = Math.max(start, limit - searchSamples); f + frame <= limit; f += frame) {
|
||||
const energy = energyAt(f, f + frame);
|
||||
if (energy <= best) {
|
||||
best = energy;
|
||||
cut = f + Math.floor(frame / 2);
|
||||
}
|
||||
}
|
||||
windows.push([start, cut]);
|
||||
start = cut;
|
||||
}
|
||||
return windows;
|
||||
}
|
||||
|
||||
/** `planChunks` over real PCM, scoring frames by mean squared amplitude. */
|
||||
export function planAudioChunks(
|
||||
audio: Float32Array,
|
||||
options: PlanChunksOptions = {},
|
||||
): Array<[number, number]> {
|
||||
return planChunks(
|
||||
audio.length,
|
||||
(start, end) => {
|
||||
let sum = 0;
|
||||
for (let i = start; i < end; i++) sum += audio[i]! * audio[i]!;
|
||||
return sum / Math.max(1, end - start);
|
||||
},
|
||||
options,
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Words → transcript chunks (the shape TranscriptView / export / live use)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** Sentence end: . ? ! optionally followed by closing quotes/brackets. */
|
||||
const SENTENCE_END = /[.?!]["'”’)\]]*$/;
|
||||
|
||||
/**
|
||||
* Group words into caption-sized chunks: a chunk ends after a sentence-final
|
||||
* word, before a pause of `gapSeconds` (default 1 s), or before a word that
|
||||
* would stretch it past `maxSegmentSeconds` (default 12 s). Chunk text has
|
||||
* Whisper's leading-space convention, so concatenating chunks yields text.
|
||||
*/
|
||||
export function wordsToChunks(
|
||||
words: readonly ParakeetWord[],
|
||||
options: { gapSeconds?: number; maxSegmentSeconds?: number } = {},
|
||||
): TranscriptChunk[] {
|
||||
const gap = options.gapSeconds ?? 1;
|
||||
const maxSegment = options.maxSegmentSeconds ?? 12;
|
||||
const chunks: TranscriptChunk[] = [];
|
||||
let current: ParakeetWord[] = [];
|
||||
|
||||
const close = () => {
|
||||
if (current.length === 0) return;
|
||||
chunks.push({
|
||||
text: ` ${current.map((w) => w.text).join(" ")}`,
|
||||
timestamp: [current[0]!.start_time, current[current.length - 1]!.end_time],
|
||||
});
|
||||
current = [];
|
||||
};
|
||||
|
||||
for (const w of words) {
|
||||
if (current.length > 0) {
|
||||
const last = current[current.length - 1]!;
|
||||
if (w.start_time - last.end_time >= gap || w.end_time - current[0]!.start_time > maxSegment) {
|
||||
close();
|
||||
}
|
||||
}
|
||||
current.push(w);
|
||||
if (SENTENCE_END.test(w.text)) close();
|
||||
}
|
||||
close();
|
||||
return chunks;
|
||||
}
|
||||
|
||||
/**
|
||||
* Run `recognize` over each planned window (a `subarray` view — never a copy
|
||||
* of the whole file) and assemble the transcript. `onWindow` fires after each
|
||||
* window with its new chunks and the absolute seconds processed so far.
|
||||
*/
|
||||
export async function transcribeWindows(
|
||||
audio: Float32Array,
|
||||
recognize: (window: Float32Array, offsetSeconds: number) => Promise<ParakeetWord[]>,
|
||||
callbacks: {
|
||||
onWindow?: (chunks: TranscriptChunk[], processedSeconds: number, totalSeconds: number) => void;
|
||||
} = {},
|
||||
): Promise<{ text: string; chunks: TranscriptChunk[] }> {
|
||||
const totalSeconds = audio.length / PARAKEET_SAMPLE_RATE;
|
||||
const chunks: TranscriptChunk[] = [];
|
||||
for (const [start, end] of planAudioChunks(audio)) {
|
||||
const words = await recognize(audio.subarray(start, end), start / PARAKEET_SAMPLE_RATE);
|
||||
const windowChunks = wordsToChunks(words);
|
||||
chunks.push(...windowChunks);
|
||||
callbacks.onWindow?.(windowChunks, end / PARAKEET_SAMPLE_RATE, totalSeconds);
|
||||
}
|
||||
return { text: chunks.map((chunk) => chunk.text).join("").trim(), chunks };
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tokenizer (vocab.txt: "<piece> <id>" per line, SentencePiece "▁" = space)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export function parseVocab(text: string): string[] {
|
||||
const id2token: string[] = [];
|
||||
for (const line of text.split(/\r?\n/)) {
|
||||
const [token, idText] = line.trim().split(/\s+/);
|
||||
const id = Number(idText);
|
||||
if (!token || !Number.isInteger(id)) continue;
|
||||
id2token[id] = token;
|
||||
}
|
||||
return id2token;
|
||||
}
|
||||
|
||||
/**
|
||||
* The tokenizer shape parakeet.js's `ParakeetModel` reads (`id2token`,
|
||||
* `blankId`, `blankToken`, `decode`). Its own `ParakeetTokenizer` is not
|
||||
* exported and only loads by URL, so this mirrors its decode rules.
|
||||
*/
|
||||
export class ParakeetVocabTokenizer {
|
||||
readonly blankToken = "<blk>";
|
||||
readonly blankId: number;
|
||||
readonly sanitizedTokens: string[];
|
||||
|
||||
constructor(readonly id2token: string[]) {
|
||||
const blank = id2token.indexOf(this.blankToken);
|
||||
this.blankId = blank === -1 ? id2token.length : blank;
|
||||
this.sanitizedTokens = id2token.map((token) => token.replace(/▁/g, " "));
|
||||
}
|
||||
|
||||
decode(ids: readonly number[]): string {
|
||||
let text = "";
|
||||
for (const id of ids) {
|
||||
if (id === this.blankId) continue;
|
||||
text += this.sanitizedTokens[id] ?? "";
|
||||
}
|
||||
return text
|
||||
.replace(/^\s+/, "")
|
||||
.replace(/\s+(?=[^\w\s])/g, "")
|
||||
.replace(/\s+/g, " ")
|
||||
.trim();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Assets
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export const PARAKEET_VOCAB_FILE = "vocab.txt";
|
||||
|
||||
/** ONNX file for one part at one dtype. Only int8 and fp16 are offered. */
|
||||
export function parakeetFileName(part: "encoder" | "decoder", dtype: WhisperDtype): string {
|
||||
const base = part === "encoder" ? "encoder-model" : "decoder_joint-model";
|
||||
if (dtype === "q8") return `${base}.int8.onnx`;
|
||||
if (dtype === "fp16") return `${base}.fp16.onnx`;
|
||||
throw new Error(`Parakeet has no ${dtype} ${part} in this app (int8 or fp16 only).`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Hub URL pinned to the measured revision. Same `/<modelId>/` shape as
|
||||
* Transformers.js downloads, so modelCache.ts's badges/storage panel/delete
|
||||
* recognise the cached files without changes (D-005).
|
||||
*/
|
||||
export function parakeetFileUrl(file: string): string {
|
||||
return `https://huggingface.co/${PARAKEET_MODEL_ID}/resolve/${PARAKEET_REVISION}/${file}`;
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
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);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user