fix(scrawl): a capture knows when each word was said

FORK.md

@@ -139,6 +139,9 @@ can be rebound in Settings → Keybindings ("Composer: Dictate").
 - It runs on the GPU when there is one (Vulkan; the log says `loaded … on Vulkan0`), about five
   times faster than the CPU. `T3CODE_DICTATION_BACKEND=cpu` keeps it on the CPU.
 - Long recordings are transcribed in pieces of at most 35 seconds, cut at pauses.
+- Canary cannot time words, and a capture needs those times to tie speech to marks. So a capture's
+  transcript comes from Parakeet (German: `parakeet-tdt-0.6b-v3`), which loads for that one run and
+  is let go; dictation stays on Canary. Without Parakeet, a capture keeps its words untimed.
 - The model (1 to 2.5 GB) loads on the first chunk of a recording and is let go after 15
   minutes without one.
 - Structuring calls `claude --print --model haiku` on the thread's server; it needs a logged-in

apps/server/src/dictation/DictationEngine.test.ts

@@ -16,7 +16,7 @@ import {
   type TranscribeModule,
 } from "./DictationEngine.ts";
 import { dictationBackend } from "./localDictation.ts";
-import { findModel, findVadModel, rankModels } from "./modelFiles.ts";
+import { findModel, findTimingModel, findVadModel, rankModels } from "./modelFiles.ts";
 import { handleDictationWorkerRequest } from "./worker.ts";
 
 const V3 = "C:/models/parakeet-tdt-0.6b-v3-Q8_0.gguf";
@@ -97,6 +97,22 @@ describe("findModel", () => {
     }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
   );
 
+  it.effect("finds only a Parakeet model that knows the language to time words", () =>
+    Effect.gen(function* () {
+      const fs = yield* FileSystem.FileSystem;
+      const path = yield* Path.Path;
+      const directory = yield* fs.makeTempDirectoryScoped();
+      const directories = [directory];
+      yield* fs.writeFileString(path.join(directory, "canary-1b-v2-Q8_0.gguf"), "");
+      expect(yield* findTimingModel({ language: "de", directories })).toBeUndefined();
+      yield* fs.writeFileString(path.join(directory, "parakeet-unified-en-0.6b-Q8_0.gguf"), "");
+      expect(yield* findTimingModel({ language: "de", directories })).toBeUndefined();
+      const v3 = path.join(directory, "parakeet-tdt-0.6b-v3-Q8_0.gguf");
+      yield* fs.writeFileString(v3, "");
+      expect(yield* findTimingModel({ language: "de", directories })).toBe(v3);
+    }).pipe(Effect.scoped, Effect.provide(NodeServices.layer)),
+  );
+
   it.effect("names the folders when nothing is installed", () =>
     Effect.gen(function* () {
       const error = yield* findModel({
@@ -168,17 +184,32 @@ class FakeModel {
   static runError: Error | undefined;
   static runLengths: Array<number> = [];
   static vadDisposed = 0;
+  static disposedPaths: Array<string> = [];
+  static timesNothing = false;
   runs: Array<{ length: number; language: string | undefined }> = [];
   disposed = false;
   readonly backend = "CPU";
-  readonly capabilities = { languages: ["de", "en"], maxTimestampKind: "word" };
+  readonly capabilities: { languages: Array<string>; maxTimestampKind: string };
+  readonly path: string;
+  constructor(path: string) {
+    this.path = path;
+    this.capabilities = {
+      languages: ["de", "en"],
+      maxTimestampKind: path.includes("canary") ? "none" : "word",
+    };
+  }
   static async load(path: string) {
     FakeModel.loads.push(path);
-    return new FakeModel();
+    return new FakeModel(path);
   }
   createSession() {
     return {
-      run: async (pcm: Float32Array, opts: { language?: string }) => {
+      run: async (pcm: Float32Array, opts: { language?: string; timestamps?: string }) => {
+        if (this.path.includes("canary")) {
+          if (opts.timestamps !== undefined) throw new Error("unsupported timestamp granularity");
+          return { text: "Canary Text", aborted: false, words: [] };
+        }
+        if (FakeModel.timesNothing) return { text: "", aborted: false, words: [] };
         this.runs.push({ length: pcm.length, language: opts.language });
         FakeModel.runLengths.push(pcm.length);
         if (FakeModel.runError !== undefined) throw FakeModel.runError;
@@ -204,6 +235,7 @@ class FakeModel {
   }
   dispose() {
     this.disposed = true;
+    FakeModel.disposedPaths.push(this.path);
   }
 }
 
@@ -220,12 +252,14 @@ const fakeModule = {
 const makeEngine = (
   resolveModel = async () => "model.gguf",
   resolveVadModel?: () => Promise<string | undefined>,
+  resolveTimingModel?: () => Promise<string | undefined>,
 ) => {
   let now = 0;
   const engine = new DictationEngine({
     loadModule: async () => fakeModule as unknown as TranscribeModule,
     resolveModel,
     ...(resolveVadModel ? { resolveVadModel } : {}),
+    ...(resolveTimingModel ? { resolveTimingModel } : {}),
     idleUnloadMs: 60_000,
     now: () => now,
   });
@@ -268,6 +302,55 @@ describe("DictationEngine", () => {
     FakeModel.vadSegments = [];
   });
 
+  describe("when the main model cannot time words", () => {
+    const canary = async () => "canary.gguf";
+    const parakeet = async () => "parakeet.gguf";
+
+    it("lets the timing model write a capture and lets go of it afterwards", async () => {
+      FakeModel.disposedPaths = [];
+      const { engine } = makeEngine(canary, undefined, parakeet);
+      engine.append("a", "de", speech(3));
+      expect(await engine.finish("a", { words: true })).toEqual({
+        text: "Hallo Welt",
+        words: [
+          { text: "Hallo", t0: 100, t1: 400 },
+          { text: "Welt", t0: 1_600, t1: 1_800 },
+        ],
+      });
+      expect(FakeModel.disposedPaths).toEqual(["parakeet.gguf"]);
+    });
+
+    it("leaves dictation with the main model", async () => {
+      FakeModel.loads = [];
+      const { engine } = makeEngine(canary, undefined, parakeet);
+      engine.append("a", "de", speech(3));
+      expect(await engine.finish("a")).toEqual({ text: "Canary Text" });
+      expect(FakeModel.loads).toEqual(["canary.gguf"]);
+    });
+
+    it("keeps the words untimed when no timing model is installed", async () => {
+      const { engine } = makeEngine(canary, undefined, async () => undefined);
+      engine.append("a", "de", speech(3));
+      expect(await engine.finish("a", { words: true })).toEqual({ text: "Canary Text" });
+    });
+
+    it("keeps the words untimed when the timing model hears nothing", async () => {
+      FakeModel.timesNothing = true;
+      const { engine } = makeEngine(canary, undefined, parakeet);
+      engine.append("a", "de", speech(3));
+      expect(await engine.finish("a", { words: true })).toEqual({ text: "Canary Text" });
+      FakeModel.timesNothing = false;
+    });
+
+    it("keeps the words untimed when the timing model cannot load", async () => {
+      const { engine } = makeEngine(canary, undefined, async () => {
+        throw new Error("broken file");
+      });
+      engine.append("a", "de", speech(3));
+      expect(await engine.finish("a", { words: true })).toEqual({ text: "Canary Text" });
+    });
+  });
+
   it("transcribes a long recording in pieces and joins them", async () => {
     const { engine } = makeEngine();
     engine.append("long", "de", speech(40));

apps/server/src/dictation/DictationEngine.ts

@@ -39,6 +39,8 @@ export interface DictationEngineOptions {
   readonly resolveModel: (language: string | undefined) => Promise<string>;
   /** Path of the Silero model, undefined when none is installed; pauses then reach the model. */
   readonly resolveVadModel?: () => Promise<string | undefined>;
+  /** Path of a model that can time words, for captures when the main model cannot. */
+  readonly resolveTimingModel?: (language: string | undefined) => Promise<string | undefined>;
   /** Let go of the model after this long without a recording; 0 keeps it. */
   readonly idleUnloadMs: number;
   /** transcribe.cpp backend; "auto" picks the GPU when there is one. */
@@ -84,6 +86,9 @@ interface RecordingSession {
   lastPreviewAt: number;
 }
 
+const modelCanTimeWords = (loaded: LoadedModel) =>
+  ["word", "token"].includes(loaded.model.capabilities.maxTimestampKind);
+
 /** "de-DE" -> "de"; the model knows languages by their two-letter code. */
 export function normalizeLanguage(language: string | undefined): string | undefined {
   const code = language?.trim().toLowerCase().split(/[-_]/)[0];
@@ -249,52 +254,109 @@ export class DictationEngine {
     // A preview run in flight is aborted instead of making the final run wait for it.
     this.#previewAbort?.abort();
     const started = this.#now();
-    const canTimeWords = ["word", "token"].includes(loaded.model.capabilities.maxTimestampKind);
-    const wantWords = options?.words === true && canTimeWords;
+    const canTimeWords = modelCanTimeWords(loaded);
+    const wantWords = options?.words === true;
     const { text, words } = await this.#exclusive(async () => {
-      const texts: Array<string> = [];
-      const words: Array<TranscriptWord> = [];
       const { speech, spans } = await this.#speechWithSpans(audio);
-      let pieceStartMs = 0;
-      for (const piece of splitAtPauses(speech)) {
-        const collect = (result: TranscribeCpp.Transcript) => {
-          texts.push(result.text);
-          if (!wantWords) return;
-          for (const word of result.words) {
-            words.push({
-              text: word.text.trim(),
-              t0: toRecordingMs(pieceStartMs + word.t0Ms, spans),
-              t1: toRecordingMs(pieceStartMs + word.t1Ms, spans),
-            });
-          }
-        };
-        try {
-          collect(
-            await loaded.session.run(piece, {
-              ...this.#runOptions(loaded, session.language),
-              ...(wantWords ? { timestamps: "word" as const } : {}),
-            }),
-          );
-        } catch (error) {
-          const module = await this.#loadModule();
-          if (error instanceof module.OutputTruncated && error.partialResult !== undefined) {
-            // One piece that ran away must not cost the rest of the recording.
-            collect(error.partialResult);
-          } else if (!(error instanceof module.InputTooShort)) {
-            throw error;
-          }
-        }
-        pieceStartMs += (piece.length / SAMPLE_RATE) * 1000;
-      }
-      return { text: texts.join(" "), words };
+      const timed =
+        wantWords && !canTimeWords
+          ? await this.#transcribeWithTimingModel(session.language, speech, spans)
+          : undefined;
+      return (
+        timed ??
+        this.#transcribeSpeech(loaded, session.language, wantWords && canTimeWords, speech, spans)
+      );
     });
     this.#lastUsedAt = this.#now();
     const cleaned = text.split(/\s+/).filter(Boolean).join(" ");
     this.#options.log?.(
       `transcribed ${(audio.length / SAMPLE_RATE).toFixed(1)}s of audio in ${this.#now() - started}ms`,
     );
     if (!cleaned) return { text: "", reason: "no-speech" };
-    return wantWords ? { text: cleaned, words } : { text: cleaned };
+    return words === undefined ? { text: cleaned } : { text: cleaned, words };
+  }
+
+  /**
+   * A capture needs to know when each word was said, which Canary cannot tell. The timing model
+   * then writes the transcript of that capture and is let go right after, so dictation keeps its
+   * own model. Undefined when it is missing or fails; the capture then keeps its words untimed.
+   */
+  async #transcribeWithTimingModel(
+    language: string | undefined,
+    speech: Float32Array,
+    spans: ReadonlyArray<KeptSpan>,
+  ): Promise<{ text: string; words: ReadonlyArray<TranscriptWord> } | undefined> {
+    try {
+      const path = await this.#options.resolveTimingModel?.(language);
+      if (path === undefined) {
+        this.#options.log?.("no model that can time words is installed, words stay untimed");
+        return undefined;
+      }
+      const module = await this.#loadModule();
+      const model = await module.TranscribeModel.load(path, {
+        backend: this.#options.backend ?? "auto",
+      });
+      const session = model.createSession();
+      try {
+        const loaded = { path, model, session, languages: model.capabilities.languages };
+        const result = await this.#transcribeSpeech(loaded, language, true, speech, spans);
+        if (result.words === undefined || result.words.length === 0) {
+          this.#options.log?.(`${path} timed no words, words stay untimed`);
+          return undefined;
+        }
+        return { text: result.text, words: result.words };
+      } finally {
+        session.dispose();
+        model.dispose();
+      }
+    } catch (error) {
+      this.#options.log?.(`timing words failed, words stay untimed: ${describe(error)}`);
+      return undefined;
+    }
+  }
+
+  async #transcribeSpeech(
+    loaded: LoadedModel,
+    language: string | undefined,
+    timed: boolean,
+    speech: Float32Array,
+    spans: ReadonlyArray<KeptSpan>,
+  ): Promise<{ text: string; words?: ReadonlyArray<TranscriptWord> }> {
+    const texts: Array<string> = [];
+    const words: Array<TranscriptWord> = [];
+    let pieceStartMs = 0;
+    for (const piece of splitAtPauses(speech)) {
+      const collect = (result: TranscribeCpp.Transcript) => {
+        texts.push(result.text);
+        if (!timed) return;
+        for (const word of result.words) {
+          words.push({
+            text: word.text.trim(),
+            t0: toRecordingMs(pieceStartMs + word.t0Ms, spans),
+            t1: toRecordingMs(pieceStartMs + word.t1Ms, spans),
+          });
+        }
+      };
+      try {
+        collect(
+          await loaded.session.run(piece, {
+            ...this.#runOptions(loaded, language),
+            ...(timed ? { timestamps: "word" as const } : {}),
+          }),
+        );
+      } catch (error) {
+        const module = await this.#loadModule();
+        if (error instanceof module.OutputTruncated && error.partialResult !== undefined) {
+          // One piece that ran away must not cost the rest of the recording.
+          collect(error.partialResult);
+        } else if (!(error instanceof module.InputTooShort)) {
+          throw error;
+        }
+      }
+      pieceStartMs += (piece.length / SAMPLE_RATE) * 1000;
+    }
+    const text = texts.join(" ");
+    return timed ? { text, words } : { text };
   }
 
   cancel(sessionId: string): void {

apps/server/src/dictation/localDictation.ts

@@ -14,7 +14,7 @@ import {
   type TranscribeModule,
   type TranscriptResult,
 } from "./DictationEngine.ts";
-import { findModel, findVadModel, modelDirectories } from "./modelFiles.ts";
+import { findModel, findTimingModel, findVadModel, modelDirectories } from "./modelFiles.ts";
 import { DEFAULT_STRUCTURE_CODEWORDS, removeFillers, splitCodeword } from "./transcriptText.ts";
 
 const IDLE_UNLOAD_MS = 15 * 60_000;
@@ -78,12 +78,20 @@ export function makeLocalDictationEngine(log: (message: string) => void): Dictat
   const resolveVadModel = Effect.fn("dictation.resolveVadModel")(function* () {
     return yield* findVadModel(modelDirectories(process.env, yield* Path.Path));
   });
+  const resolveTimingModel = Effect.fn("dictation.resolveTimingModel")(function* (
+    language: string | undefined,
+  ) {
+    const directories = modelDirectories(process.env, yield* Path.Path);
+    return yield* findTimingModel({ language, directories });
+  });
   return new DictationEngine({
     loadModule: loadTranscribeCpp,
     resolveModel: (language) =>
       Effect.runPromise(resolveModel(language).pipe(Effect.provide(NodeServices.layer))),
     resolveVadModel: () =>
       Effect.runPromise(resolveVadModel().pipe(Effect.provide(NodeServices.layer))),
+    resolveTimingModel: (language) =>
+      Effect.runPromise(resolveTimingModel(language).pipe(Effect.provide(NodeServices.layer))),
     idleUnloadMs: IDLE_UNLOAD_MS,
     backend: dictationBackend(process.env.T3CODE_DICTATION_BACKEND),
     log,