@drakulavich/parakeet-cli 0.1.4 → 0.2.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@drakulavich/parakeet-cli",
-  "version": "0.1.4",
+  "version": "0.2.0",
   "description": "Fast multilingual speech-to-text CLI powered by NVIDIA Parakeet ONNX models",
   "type": "module",
   "bin": {

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/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/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";
@@ -63,6 +63,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);
   return tokenizer.detokenize(tokens);
 }