diff --git a/src/lib/parakeet.ts b/src/lib/parakeet.ts new file mode 100644 index 0000000..51fb5aa --- /dev/null +++ b/src/lib/parakeet.ts @@ -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, + 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: " " 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 = ""; + 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 `//` 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}`; +} diff --git a/tests/parakeet.test.ts b/tests/parakeet.test.ts new file mode 100644 index 0000000..62a42b5 --- /dev/null +++ b/tests/parakeet.test.ts @@ -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 = " 0\n▁the 1\n▁cat 2\ns 3\n. 4\n 5\n"; + + test("parseVocab indexes tokens by id", () => { + expect(parseVocab(vocab)).toEqual(["", "▁the", "▁cat", "s", ".", ""]); + }); + + test("decode joins pieces, skips blank, and attaches punctuation", () => { + const tokenizer = new ParakeetVocabTokenizer(parseVocab(vocab)); + expect(tokenizer.blankId).toBe(5); + expect(tokenizer.blankToken).toBe(""); + 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); + }); +});