@drakulavich/parakeet-cli 0.1.4 → 0.3.0

package/package.json

@@ -1,11 +1,15 @@
 {
   "name": "@drakulavich/parakeet-cli",
-  "version": "0.1.4",
+  "version": "0.3.0",
   "description": "Fast multilingual speech-to-text CLI powered by NVIDIA Parakeet ONNX models",
   "type": "module",
   "bin": {
     "parakeet": "bin/parakeet.js"
   },
+  "exports": {
+    ".": "./src/cli.ts",
+    "./core": "./src/lib.ts"
+  },
   "files": [
     "bin/",
     "src/",

package/src/__tests__/decoder.test.ts

@@ -1,5 +1,5 @@
 import { describe, test, expect } from "bun:test";
-import { greedyDecode, type DecoderSession } from "../decoder";
+import { beamDecode, type DecoderSession } from "../decoder";
 
 function mockSession(responses: Array<{ tokenLogits: number[]; durationLogits: number[] }>): DecoderSession {
   let callIndex = 0;
@@ -24,7 +24,8 @@ describe("decoder", () => {
       { tokenLogits: [0, 10, 0, -10], durationLogits: [10, 0] },
       { tokenLogits: [0, 0, 0, 10], durationLogits: [10, 0] },
     ]);
-    const tokens = await greedyDecode(session, 3);
+    const encoderData = new Float32Array(3);
+    const tokens = await beamDecode(session, 3, encoderData, 1, 1);
     expect(tokens).toEqual([0, 1]);
   });
 
@@ -34,24 +35,16 @@ describe("decoder", () => {
       { tokenLogits: [0, 10, 0, -10], durationLogits: [10, 0, 0] },
       { tokenLogits: [0, 0, 0, 10], durationLogits: [10, 0, 0] },
     ]);
-    const tokens = await greedyDecode(session, 5);
+    const encoderData = new Float32Array(5);
+    const tokens = await beamDecode(session, 5, encoderData, 1, 1);
     expect(tokens).toEqual([0, 1]);
   });
 
-  test("handles max_tokens_per_step limit", async () => {
-    const session = mockSession([
-      { tokenLogits: [10, 0, 0, -10], durationLogits: [10, 0] },
-    ]);
-    const tokens = await greedyDecode(session, 2);
-    expect(tokens.length).toBeLessThanOrEqual(20);
-    expect(tokens.length).toBeGreaterThan(0);
-  });
-
   test("returns empty for zero-length encoder output", async () => {
     const session = mockSession([
       { tokenLogits: [0, 0, 0, 10], durationLogits: [10, 0] },
     ]);
-    const tokens = await greedyDecode(session, 0);
+    const tokens = await beamDecode(session, 0, new Float32Array(0), 1);
     expect(tokens).toEqual([]);
   });
 });

package/src/__tests__/lib.test.ts

@@ -0,0 +1,12 @@
+import { describe, expect, it } from "bun:test";
+import { transcribe } from "../lib";
+
+describe("lib API", () => {
+  it("exports transcribe function", () => {
+    expect(typeof transcribe).toBe("function");
+  });
+
+  it("rejects missing file", async () => {
+    await expect(transcribe("/nonexistent/audio.wav")).rejects.toThrow("File not found");
+  });
+});

package/src/cli.ts

@@ -1,7 +1,6 @@
 #!/usr/bin/env bun
 
-import { existsSync } from "fs";
-import { transcribe } from "./transcribe";
+import { transcribe } from "./lib";
 
 async function main(): Promise<void> {
   const args = process.argv.slice(2);
@@ -20,11 +19,6 @@ async function main(): Promise<void> {
     process.exit(1);
   }
 
-  if (!existsSync(file)) {
-    console.error(`Error: file not found: ${file}`);
-    process.exit(1);
-  }
-
   try {
     const text = await transcribe(file, { noCache });
     if (text) process.stdout.write(text + "\n");

package/src/decoder.ts

@@ -20,69 +20,87 @@ export interface DecoderSession {
   stateDims: { layers: number; hidden: number };
 }
 
-export async function greedyDecode(
+const DEFAULT_BEAM_WIDTH = 4;
+
+interface Beam {
+  tokens: number[];
+  score: number;
+  lastToken: number;
+  state1: F32;
+  state2: F32;
+  t: number;
+}
+
+export async function beamDecode(
   session: DecoderSession,
   encoderLength: number,
-  encoderData?: Float32Array,
-  encoderDim?: number
+  encoderData: Float32Array,
+  encoderDim: number,
+  beamWidth: number = DEFAULT_BEAM_WIDTH,
 ): Promise<number[]> {
   if (encoderLength === 0) return [];
 
-  const tokens: number[] = [];
   const stateSize = session.stateDims.layers * session.stateDims.hidden;
-  let state1: F32 = new Float32Array(stateSize);
-  let state2: F32 = new Float32Array(stateSize);
-  let lastToken = session.blankId;
-
-  let t = 0;
-  while (t < encoderLength) {
-    let tokensThisStep = 0;
-
-    while (tokensThisStep < MAX_TOKENS_PER_STEP) {
-      let frame: Float32Array;
-      if (encoderData && encoderDim) {
-        // Must copy — ort.Tensor doesn't work with subarray views under Bun
-        frame = encoderData.slice(t * encoderDim, (t + 1) * encoderDim);
-      } else {
-        frame = new Float32Array(1);
-      }
 
-      const result = await session.decode(frame, [lastToken], 1, state1, state2);
+  let beams: Beam[] = [{
+    tokens: [],
+    score: 0,
+    lastToken: session.blankId,
+    state1: new Float32Array(stateSize),
+    state2: new Float32Array(stateSize),
+    t: 0,
+  }];
+
+  const maxSteps = encoderLength * MAX_TOKENS_PER_STEP;
+
+  for (let step = 0; step < maxSteps; step++) {
+    const active = beams.filter(b => b.t < encoderLength);
+    if (active.length === 0) break;
+
+    const candidates: Beam[] = [];
+
+    for (const beam of active) {
+      // Must copy — ort.Tensor doesn't work with subarray views under Bun
+      const frame = encoderData.slice(beam.t * encoderDim, (beam.t + 1) * encoderDim);
+      const result = await session.decode(frame, [beam.lastToken], 1, beam.state1, beam.state2);
       const output = result.output;
 
       const tokenLogits = output.slice(0, session.vocabSize);
       const durationLogits = output.slice(session.vocabSize);
-
-      const tokenId = argmax(tokenLogits);
       const duration = argmax(durationLogits);
 
-      state1 = result.state1;
-      state2 = result.state2;
-
-      if (tokenId === session.blankId) {
-        t += 1;
-        break;
-      }
-
-      tokens.push(tokenId);
-      lastToken = tokenId;
-      tokensThisStep++;
+      // Blank option: advance one frame, keep same tokens
+      candidates.push({
+        tokens: beam.tokens,
+        score: beam.score + tokenLogits[session.blankId],
+        lastToken: beam.lastToken,
+        state1: result.state1,
+        state2: result.state2,
+        t: beam.t + 1,
+      });
 
-      if (duration > 0) {
-        t += duration;
-        break;
+      // Top non-blank token options
+      const topK = topKIndices(tokenLogits, beamWidth, session.blankId);
+      for (const tokenId of topK) {
+        candidates.push({
+          tokens: [...beam.tokens, tokenId],
+          score: beam.score + tokenLogits[tokenId],
+          lastToken: tokenId,
+          state1: result.state1,
+          state2: result.state2,
+          t: duration > 0 ? beam.t + duration : beam.t,
+        });
       }
     }
 
-    if (tokensThisStep >= MAX_TOKENS_PER_STEP) {
-      t += 1;
-    }
+    candidates.sort((a, b) => b.score - a.score);
+    beams = candidates.slice(0, beamWidth);
   }
 
-  return tokens;
+  return beams[0].tokens;
 }
 
-function argmax(arr: Float32Array): number {
+function argmax(arr: F32): number {
   let maxIdx = 0;
   let maxVal = arr[0];
   for (let i = 1; i < arr.length; i++) {
@@ -94,6 +112,15 @@ function argmax(arr: Float32Array): number {
   return maxIdx;
 }
 
+function topKIndices(arr: F32, k: number, excludeId: number): number[] {
+  const indexed: [number, number][] = [];
+  for (let i = 0; i < arr.length; i++) {
+    if (i !== excludeId) indexed.push([arr[i], i]);
+  }
+  indexed.sort((a, b) => b[0] - a[0]);
+  return indexed.slice(0, k).map(([, i]) => i);
+}
+
 let onnxSession: ort.InferenceSession | null = null;
 
 export async function initDecoder(modelDir: string): Promise<void> {

package/src/lib.ts

@@ -0,0 +1,23 @@
+import { existsSync } from "fs";
+import { transcribe as internalTranscribe } from "./transcribe";
+
+export interface TranscribeOptions {
+  beamWidth?: number;
+  noCache?: boolean;
+  modelDir?: string;
+}
+
+export async function transcribe(
+  audioPath: string,
+  options: TranscribeOptions = {},
+): Promise<string> {
+  if (!existsSync(audioPath)) {
+    throw new Error(`File not found: ${audioPath}`);
+  }
+
+  return internalTranscribe(audioPath, {
+    noCache: options.noCache ?? false,
+    beamWidth: options.beamWidth,
+    modelDir: options.modelDir,
+  });
+}

package/src/models.ts

@@ -16,13 +16,13 @@ export function getModelDir(): string {
   return join(homedir(), ".cache", "parakeet", "v3");
 }
 
-export function isModelCached(): boolean {
-  const dir = getModelDir();
-  return MODEL_FILES.every((f) => existsSync(join(dir, f)));
+export function isModelCached(dir?: string): boolean {
+  const d = dir ?? getModelDir();
+  return MODEL_FILES.every((f) => existsSync(join(d, f)));
 }
 
-export async function ensureModel(noCache = false): Promise<string> {
-  const dir = getModelDir();
+export async function ensureModel(noCache = false, modelDir?: string): Promise<string> {
+  const dir = modelDir ?? getModelDir();
 
   if (!noCache && isModelCached()) {
     return dir;

package/src/transcribe.ts

@@ -5,7 +5,7 @@ import { initEncoder, encode } from "./encoder";
 import {
   initDecoder,
   createOnnxDecoderSession,
-  greedyDecode,
+  beamDecode,
 } from "./decoder";
 import { Tokenizer } from "./tokenizer";
 import { join } from "path";
@@ -26,6 +26,8 @@ const DECODER_HIDDEN = 640;
 
 export interface TranscribeOptions {
   noCache?: boolean;
+  beamWidth?: number;
+  modelDir?: string;
 }
 
 // Minimum 0.1s of audio at 16kHz to produce meaningful output
@@ -39,7 +41,8 @@ export async function transcribe(audioPath: string, opts: TranscribeOptions = {}
   }
 
   const noCache = opts.noCache ?? false;
-  const modelDir = await ensureModel(noCache);
+  const beamWidth = opts.beamWidth ?? 4;
+  const modelDir = await ensureModel(noCache, opts.modelDir);
   const tokenizer = await Tokenizer.fromFile(join(modelDir, "vocab.txt"));
 
   await initPreprocessor(modelDir);
@@ -63,6 +66,6 @@ export async function transcribe(audioPath: string, opts: TranscribeOptions = {}
     DECODER_HIDDEN,
   );
 
-  const tokens = await greedyDecode(session, encodedLength, transposed, D);
+  const tokens = await beamDecode(session, encodedLength, transposed, D, beamWidth);
   return tokenizer.detokenize(tokens);
 }