@tychenjiajun/pi-voice 0.7.0 → 0.8.1

package/src/transcription.ts

@@ -1,22 +1,28 @@
 import type {
+  BackendInfo,
   Capabilities,
+  LangIdResult,
   Session,
   Stream,
   TranscribeModel,
 } from "transcribe-cpp";
-import { convertChineseOutput, isChineseLanguage } from "./chinese.js";
-import { writeLog, type LogLevel } from "./log.js";
+import { convertChineseOutput, shouldConvertChinese } from "./chinese.js";
+import { describeError, log, writeLog, type LogLevel } from "./log.js";
 import type { ChineseOutput } from "./settings.js";
 
 export type TranscriptionOptions = {
   signal?: AbortSignal;
   language?: string;
   chineseOutput?: ChineseOutput;
+  /** Chinese is a preferred or chosen language: convert any Chinese text. */
+  chineseEnabled?: boolean;
 };
 
 export type DictationStream = {
-  feed(chunk: Float32Array): Promise<void>;
-  finalize(): Promise<string>;
+  /** `probs` are the chunk's Silero frame probabilities, when the take has them. */
+  feed(chunk: Float32Array, probs?: readonly number[]): Promise<void>;
+  /** `audio` is the whole take, for a stream that checks it saw all of it. */
+  finalize(audio?: Float32Array): Promise<string>;
   reset(): void;
 };
 
@@ -33,18 +39,18 @@ function validateLanguage(
 
 /**
  * Shared final step for batch and streaming output: trim, then apply the
- * Chinese script preference when the detected (or configured) language is
- * Chinese. Keeping this in one place keeps the two paths from drifting.
+ * Chinese script preference (see shouldConvertChinese). Keeping this in one
+ * place keeps the two paths from drifting.
  */
 async function finishTranscript(
   text: string,
   detectedLanguage: string,
-  configuredLanguage: string | undefined,
-  chineseOutput: ChineseOutput,
+  options: TranscriptionOptions,
 ): Promise<string> {
   const trimmed = text.trim();
-  return isChineseLanguage(detectedLanguage || configuredLanguage || "")
-    ? convertChineseOutput(trimmed, chineseOutput)
+  const language = detectedLanguage || options.language || "";
+  return shouldConvertChinese(trimmed, { chineseEnabled: options.chineseEnabled ?? false, language })
+    ? convertChineseOutput(trimmed, options.chineseOutput ?? "simplified")
     : trimmed;
 }
 
@@ -54,8 +60,7 @@ class TranscribeCppDictationStream implements DictationStream {
   constructor(
     private readonly session: Session,
     private readonly stream: Stream,
-    private readonly language: string | undefined,
-    private readonly chineseOutput: ChineseOutput,
+    private readonly options: TranscriptionOptions,
   ) {}
 
   async feed(chunk: Float32Array): Promise<void> {
@@ -80,7 +85,7 @@ class TranscribeCppDictationStream implements DictationStream {
       this.close();
     }
 
-    return finishTranscript(text, detectedLanguage, this.language, this.chineseOutput);
+    return finishTranscript(text, detectedLanguage, this.options);
   }
 
   reset(): void {
@@ -103,30 +108,60 @@ const TRANSCRIBE_LEVELS: Record<number, LogLevel> = { 1: "debug", 2: "warn", 3:
 const TRANSCRIBE_CONTINUATION = 5;
 let lastTranscribeLevel: LogLevel = "debug";
 
-function routeTranscribeLog(setLogHandler: typeof import("transcribe-cpp").setLogHandler): void {
+export function routeTranscribeLog(setLogHandler: typeof import("transcribe-cpp").setLogHandler): void {
   setLogHandler((level, message) => {
     if (level !== TRANSCRIBE_CONTINUATION) lastTranscribeLevel = TRANSCRIBE_LEVELS[level] ?? "debug";
     writeLog(lastTranscribeLevel, "transcribe", message);
   });
 }
 
-/** A reusable loaded transcribe.cpp model. Calls must be scheduled sequentially. */
+type TranscribeCppRuntime = Pick<
+  typeof import("transcribe-cpp"),
+  "getAvailableBackends" | "OutOfMemory" | "setLogHandler" | "TranscribeModel"
+>;
+
+type RuntimeLoader = () => Promise<TranscribeCppRuntime>;
+
+// Best-effort: load failures don't report their device, so any non-CPU device is assumed to have been used.
+function cpuFallbackDevice(runtime: TranscribeCppRuntime): BackendInfo | undefined {
+  const devices = runtime.getAvailableBackends();
+  const acceleratorAvailable = devices.some((device) => device.deviceType !== "cpu");
+  return acceleratorAvailable
+    ? devices.find((device) => device.deviceType === "cpu")
+    : undefined;
+}
+
+/**
+ * A reusable loaded transcribe.cpp model: an ASR model, or the language
+ * detector. Calls must be scheduled sequentially.
+ */
 export class TranscribeCppBackend {
   private model: TranscribeModel | undefined;
   private loading: Promise<TranscribeModel> | undefined;
   private disposed = false;
 
-  constructor(private readonly modelPath: string) {}
+  constructor(
+    private readonly modelPath: string,
+    private readonly loadRuntime: RuntimeLoader = () => import("transcribe-cpp"),
+  ) {}
 
   async prepare(): Promise<void> {
     if (this.model) return;
     if (this.disposed) throw new Error("Transcription backend has been disposed");
 
     if (!this.loading) {
-      this.loading = import("transcribe-cpp")
-        .then(({ TranscribeModel, setLogHandler }) => {
-          routeTranscribeLog(setLogHandler);
-          return TranscribeModel.load(this.modelPath);
+      this.loading = this.loadRuntime()
+        .then(async (runtime) => {
+          routeTranscribeLog(runtime.setLogHandler);
+          try {
+            return await runtime.TranscribeModel.load(this.modelPath);
+          } catch (error) {
+            if (!(error instanceof runtime.OutOfMemory)) throw error;
+            const cpu = cpuFallbackDevice(runtime);
+            if (!cpu) throw error;
+            log.warn(`accelerator model load ran out of memory; retrying on CPU: ${describeError(error)}`);
+            return runtime.TranscribeModel.load(this.modelPath, { device: cpu });
+          }
         })
         .then((model) => {
           if (this.disposed) {
@@ -163,8 +198,7 @@ export class TranscribeCppBackend {
       return new TranscribeCppDictationStream(
         session,
         stream,
-        options.language,
-        options.chineseOutput ?? "simplified",
+        options,
       );
     } catch (error) {
       session.dispose();
@@ -189,11 +223,25 @@ export class TranscribeCppBackend {
     return finishTranscript(
       result.text,
       result.language,
-      options.language,
-      options.chineseOutput ?? "simplified",
+      options,
     );
   }
 
+  /** Rank `allowed` detector labels for one clip (a "langid" model only). */
+  async identifyLanguage(
+    pcm: Float32Array,
+    allowed: readonly string[],
+    signal?: AbortSignal,
+  ): Promise<LangIdResult> {
+    await this.prepare();
+    const session = this.model!.createLangIdSession();
+    try {
+      return await session.run(pcm, { allowed, signal });
+    } finally {
+      session.dispose();
+    }
+  }
+
   async dispose(): Promise<void> {
     this.disposed = true;
     await this.loading?.catch(() => undefined);

package/src/try-it.ts

@@ -18,6 +18,7 @@ import type { TranscribeSettings } from "./settings.js";
 import { displayShortcut } from "./shortcut-core.js";
 import { TranscriptionService } from "./transcription-service.js";
 import { TranscriptPreview } from "./transcript-preview.js";
+import { VoiceActivityDetection } from "./vad.js";
 import { editorBorder, onboardingHeader, PANEL_PADDING, panelBorder, paneRowBudget } from "./ui-components.js";
 import {
   formatTranscriptionSummary,
@@ -28,7 +29,7 @@ import {
 
 type UiTheme = ExtensionContext["ui"]["theme"];
 
-type TryItPaneOptions = Pick<DictationControllerOptions, "openCapture" | "now"> & {
+type TryItPaneOptions = Pick<DictationControllerOptions, "openCapture" | "now" | "createSpeechFilter"> & {
   /** Shown only before the first recording attempt when macOS has not asked yet. */
   showMacPermissionNote?: boolean;
 };
@@ -83,6 +84,7 @@ export class TryItPane implements Component {
     this.dictation = new DictationController(service, {
       openCapture: options.openCapture,
       now: options.now,
+      createSpeechFilter: options.createSpeechFilter,
       onChange: () => this.refresh(),
       onAudio: (chunk) => {
         this.analyzer.push(chunk);
@@ -154,6 +156,10 @@ export class TryItPane implements Component {
       } else if (localModel && needsFasterModel(speechSeconds, transcribeSeconds)) {
         details = fg("warning", `Slow on this machine? Press ${this.keys.keyText("voice.tryIt.model")} to try another model.`);
       }
+      if (state.result.warning) {
+        const warning = fg("warning", state.result.warning);
+        details = details ? `${details}\n${warning}` : warning;
+      }
     } else if (state.phase === "error") {
       activity = fg("error", state.stage === "model" ? "Could not load the model" : state.stage === "capture" ? "Microphone capture failed" : "Transcription failed");
       const message = state.cause instanceof Error ? state.cause.message : String(state.cause);
@@ -311,17 +317,20 @@ export class TryItPane implements Component {
 }
 
 export async function tryVoice(ctx: ExtensionContext, settings: TranscribeSettings): Promise<TryItResult | undefined> {
-  const service = new TranscriptionService();
+  const voiceActivity = new VoiceActivityDetection();
+  const service = new TranscriptionService(undefined, (pcm, signal) => voiceActivity.scoreClip(pcm, signal));
   let pane: TryItPane | undefined;
   try {
     return await ctx.ui.custom<TryItResult>((tui, theme, keybindings, done) =>
       (pane = new TryItPane(tui, theme, keybindings, settings, service, done, {
         openCapture: openMicrophone,
+        createSpeechFilter: (configured, onSpeech) => voiceActivity.createFilter(configured, onSpeech),
         showMacPermissionNote: process.platform === "darwin" && microphonePermission() === "not-determined",
       })),
     );
   } finally {
     await pane?.dispose();
     await service.shutdown();
+    await voiceActivity.dispose();
   }
 }

package/src/vad.ts

@@ -0,0 +1,301 @@
+import { fileURLToPath } from "node:url";
+import type { TranscribeModel } from "transcribe-cpp";
+import { CAPTURE_CHUNK_SAMPLES, CAPTURE_SAMPLE_RATE } from "./audio-constants.js";
+import { getCatalogModel } from "./catalog.js";
+import { chunksTranscription } from "./chunking.js";
+import type { OnSpeech, SpeechFilter } from "./dictation-controller.js";
+import { describeError, log } from "./log.js";
+import type { TranscribeSettings } from "./settings.js";
+
+// Constants come from Handy (cjpais/handy, src-tauri/src/audio_toolkit/vad).
+// It's SmoothedVad idea to improve basic VAD performance
+const THRESHOLD = 0.3;
+const PREFILL_MS = 450;
+const ONSET_MS = 60;
+const OFFLINE_HANGOVER_MS = 450;
+const STREAMING_HANGOVER_MS = 1650;
+const MODEL_PATH = fileURLToPath(new URL("../models/silero-vad-v6.2-F32.gguf", import.meta.url));
+
+function framesFor(durationMs: number): number {
+  return Math.ceil((durationMs * CAPTURE_SAMPLE_RATE) / (CAPTURE_CHUNK_SAMPLES * 1000));
+}
+
+function concat(pieces: readonly Float32Array[]): Float32Array {
+  if (pieces.length === 1) return pieces[0]!;
+  const out = new Float32Array(pieces.reduce((sum, piece) => sum + piece.length, 0));
+  let offset = 0;
+  for (const piece of pieces) {
+    out.set(piece, offset);
+    offset += piece.length;
+  }
+  return out;
+}
+
+type BufferedFrame<T> = { frame: T; emitted: boolean };
+
+// Handy's SmoothedVad
+export class SpeechGate<T = Float32Array> {
+  private readonly buffer: BufferedFrame<T>[] = [];
+  private hangover = 0;
+  private onset = 0;
+  private speaking = false;
+
+  constructor(private readonly options: { prefillFrames: number; onsetFrames: number; hangoverFrames: number }) {}
+
+  get inSpeech(): boolean { return this.speaking; }
+
+  /** The frames to keep now, in order. */
+  push(frame: T, voiced: boolean): T[] {
+    const buffered: BufferedFrame<T> = { frame, emitted: false };
+    this.buffer.push(buffered);
+    while (this.buffer.length > this.options.prefillFrames + 1) this.buffer.shift();
+
+    if (!this.speaking) {
+      if (!voiced) {
+        this.onset = 0;
+        return [];
+      }
+      if (++this.onset < this.options.onsetFrames) return [];
+      this.speaking = true;
+      this.onset = 0;
+      this.hangover = this.options.hangoverFrames;
+      return this.takeWithheld();
+    }
+
+    if (voiced) this.hangover = this.options.hangoverFrames;
+    else if (this.hangover > 0) this.hangover--;
+    else {
+      this.speaking = false;
+      return [];
+    }
+    buffered.emitted = true;
+    return [frame];
+  }
+
+  takeWithheld(): T[] {
+    const out: T[] = [];
+    for (const buffered of this.buffer) {
+      if (buffered.emitted) continue;
+      buffered.emitted = true;
+      out.push(buffered.frame);
+    }
+    return out;
+  }
+
+  finish(): T[] {
+    return this.onset > 0 ? this.takeWithheld() : [];
+  }
+}
+
+export type VadScorer = {
+  score(frame: Float32Array): Promise<number>;
+  dispose(): void;
+};
+
+type VoiceActivityFilterOptions = {
+  /** Opened on the first audio. */
+  openScorer: () => Promise<VadScorer>;
+  /** Receives kept audio, in order, as soon as it is decided, with its frame probabilities. */
+  onSpeech: OnSpeech;
+  hangoverMs: number;
+  /** False keeps every frame: the probabilities are still wanted, to cut chunks. */
+  dropSilence?: boolean;
+};
+
+/** A frame and its speech probability; undefined once the detector failed. */
+type ScoredFrame = { samples: Float32Array; prob?: number };
+
+/**
+ * One recording's VAD. Each pushed chunk is one frame; scoring is async, so
+ * frames queue and are decided in order. If the detector fails, nothing not
+ * yet dropped is lost and the rest of the recording passes through.
+ */
+export class VoiceActivityFilter implements SpeechFilter {
+  private readonly gate: SpeechGate<ScoredFrame>;
+  private readonly dropSilence: boolean;
+  private readonly pending: Float32Array[] = [];
+  private readonly kept: Float32Array[] = [];
+  private recordedSamples = 0;
+  private scorer: VadScorer | undefined;
+  private running: Promise<void> | undefined;
+  private failed = false;
+  private finishing = false;
+  private discarded = false;
+
+  constructor(private readonly options: VoiceActivityFilterOptions) {
+    this.dropSilence = options.dropSilence ?? true;
+    this.gate = new SpeechGate({
+      prefillFrames: framesFor(PREFILL_MS),
+      onsetFrames: framesFor(ONSET_MS),
+      hangoverFrames: framesFor(options.hangoverMs),
+    });
+  }
+
+  push(chunk: Float32Array): void {
+    if (this.finishing || this.discarded) return;
+    this.pending.push(chunk);
+    this.kick();
+  }
+
+  async finish(): Promise<Float32Array> {
+    this.finishing = true;
+    while (this.running) await this.running;
+    if (this.discarded) return new Float32Array();
+    if (!this.failed) this.keep(this.gate.finish());
+    this.release();
+    const kept = concat(this.kept);
+    log.info(this.failed
+      ? "voice activity detection unavailable; kept the whole recording"
+      : `voice activity: kept ${(kept.length / CAPTURE_SAMPLE_RATE).toFixed(1)} s of ${(this.recordedSamples / CAPTURE_SAMPLE_RATE).toFixed(1)} s`);
+    return kept;
+  }
+
+  discard(): void {
+    if (this.discarded) return;
+    this.discarded = true;
+    this.pending.length = 0;
+    // A score may be in flight; free the session once the loop has stopped.
+    if (this.running) void this.running.then(() => this.release());
+    else this.release();
+  }
+
+  private kick(): void {
+    if (this.running || this.discarded || this.pending.length === 0) return;
+    const running = this.drain().catch((error: unknown) => this.fail(error)).finally(() => {
+      if (this.running === running) this.running = undefined;
+      this.kick();
+    });
+    this.running = running;
+  }
+
+  private async drain(): Promise<void> {
+    if (!this.scorer && !this.failed) {
+      try {
+        this.scorer = await this.options.openScorer();
+      } catch (error) {
+        this.fail(error);
+      }
+      if (this.discarded) return;
+    }
+    while (this.pending.length > 0 && !this.discarded) {
+      const frame = this.pending.shift()!;
+      this.recordedSamples += frame.length;
+      if (this.failed) {
+        this.keep([{ samples: frame }]);
+        continue;
+      }
+      // Only a recording's last chunk is short; it follows the current
+      // decision and has no probability of its own.
+      if (frame.length < CAPTURE_CHUNK_SAMPLES) {
+        if (this.gate.inSpeech || !this.dropSilence) this.keep([{ samples: frame }]);
+        continue;
+      }
+      let prob: number;
+      try {
+        prob = await this.scorer!.score(frame);
+      } catch (error) {
+        this.fail(error);
+        this.keep([{ samples: frame }]);
+        continue;
+      }
+      if (this.discarded) return;
+      const scored = { samples: frame, prob };
+      this.keep(this.dropSilence ? this.gate.push(scored, prob > THRESHOLD) : [scored]);
+    }
+  }
+
+  private fail(error: unknown): void {
+    if (this.failed) return;
+    this.failed = true;
+    log.warn(`voice activity detection unavailable, keeping all audio: ${describeError(error)}`);
+    this.keep(this.gate.takeWithheld());
+    this.release();
+  }
+
+  private keep(frames: ScoredFrame[]): void {
+    if (frames.length === 0 || this.discarded) return;
+    const chunk = concat(frames.map((frame) => frame.samples));
+    this.kept.push(chunk);
+    const whole = frames.filter((frame) => frame.samples.length === CAPTURE_CHUNK_SAMPLES);
+    const probs = whole.every((frame) => frame.prob !== undefined) ? whole.map((frame) => frame.prob!) : undefined;
+    try { this.options.onSpeech(chunk, probs); } catch { /* Still in the final transcript. */ }
+  }
+
+  private release(): void {
+    this.scorer?.dispose();
+    this.scorer = undefined;
+  }
+}
+
+/** Loads the bundled Silero model on first use and keeps it across recordings. */
+export class VoiceActivityDetection {
+  private model: Promise<TranscribeModel> | undefined;
+  private disposed = false;
+
+  /**
+   * Undefined when the setting is off, unless the model is chunked: then
+   * every frame is kept, but scored for the chunk cuts.
+   */
+  createFilter(settings: TranscribeSettings, onSpeech: OnSpeech): SpeechFilter | undefined {
+    const chunked = chunksTranscription(settings);
+    if ((!settings.voiceActivityDetection && !chunked) || this.disposed) return undefined;
+    const streaming = settings.model
+      ? getCatalogModel(settings.model.id)?.capabilities.streaming ?? false
+      : false;
+    return new VoiceActivityFilter({
+      openScorer: () => this.openScorer(),
+      onSpeech,
+      hangoverMs: streaming ? STREAMING_HANGOVER_MS : OFFLINE_HANGOVER_MS,
+      dropSilence: settings.voiceActivityDetection,
+    });
+  }
+
+  /** Silero's probability for each 512-sample frame of a whole clip (the last zero padded). */
+  async scoreClip(pcm: Float32Array, signal?: AbortSignal): Promise<number[]> {
+    const session = (await this.loaded()).createVadSession();
+    try {
+      return (await session.run(pcm, { signal })).probs;
+    } finally {
+      session.dispose();
+    }
+  }
+
+  private async loaded(): Promise<TranscribeModel> {
+    if (this.disposed) throw new Error("Pi Voice is shutting down");
+    // A failed load is retried by the next use.
+    const loading = (this.model ??= loadModel().catch((error: unknown) => {
+      this.model = undefined;
+      throw error;
+    }));
+    const model = await loading;
+    if (this.disposed) throw new Error("Pi Voice is shutting down");
+    return model;
+  }
+
+  private async openScorer(): Promise<VadScorer> {
+    const session = (await this.loaded()).createVadSession({ nThreads: 1 });
+    return {
+      async score(frame) {
+        const { probs } = await session.streamFeed(frame);
+        if (probs.length !== 1) throw new Error(`Silero VAD returned ${probs.length} probabilities for one frame`);
+        return probs[0]!;
+      },
+      dispose: () => session.dispose(),
+    };
+  }
+
+  async dispose(): Promise<void> {
+    this.disposed = true;
+    const model = await this.model?.catch(() => undefined);
+    this.model = undefined;
+    model?.dispose();
+  }
+}
+
+async function loadModel(): Promise<TranscribeModel> {
+  // Route native logs into Pi Voice's log before the first native call.
+  const { routeTranscribeLog } = await import("./transcription.js");
+  const { TranscribeModel, setLogHandler } = await import("transcribe-cpp");
+  routeTranscribeLog(setLogHandler);
+  return TranscribeModel.load(MODEL_PATH);
+}