diff --git a/FORK.md b/FORK.md index 8265ea77cf9a7..1eeaf522adb89 100644 --- a/FORK.md +++ b/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 diff --git a/apps/server/src/dictation/DictationEngine.test.ts b/apps/server/src/dictation/DictationEngine.test.ts index 960a40f3242c1..1494dc6681de5 100644 --- a/apps/server/src/dictation/DictationEngine.test.ts +++ b/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 = []; static vadDisposed = 0; + static disposedPaths: Array = []; + 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; 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, + resolveTimingModel?: () => Promise, ) => { 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)); diff --git a/apps/server/src/dictation/DictationEngine.ts b/apps/server/src/dictation/DictationEngine.ts index 5873664aa59bc..c422949cd46c1 100644 --- a/apps/server/src/dictation/DictationEngine.ts +++ b/apps/server/src/dictation/DictationEngine.ts @@ -39,6 +39,8 @@ export interface DictationEngineOptions { readonly resolveModel: (language: string | undefined) => Promise; /** Path of the Silero model, undefined when none is installed; pauses then reach the model. */ readonly resolveVadModel?: () => Promise; + /** Path of a model that can time words, for captures when the main model cannot. */ + readonly resolveTimingModel?: (language: string | undefined) => Promise; /** 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,44 +254,18 @@ 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 = []; - const words: Array = []; 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(" "); @@ -294,7 +273,90 @@ export class DictationEngine { `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, + ): Promise<{ text: string; words: ReadonlyArray } | 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, + ): Promise<{ text: string; words?: ReadonlyArray }> { + const texts: Array = []; + const words: Array = []; + 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 { diff --git a/apps/server/src/dictation/localDictation.ts b/apps/server/src/dictation/localDictation.ts index d51c67d831464..a31f42e78d20c 100644 --- a/apps/server/src/dictation/localDictation.ts +++ b/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, diff --git a/apps/server/src/dictation/modelFiles.ts b/apps/server/src/dictation/modelFiles.ts index 79f1dc02fc698..3144a4aa43e8f 100644 --- a/apps/server/src/dictation/modelFiles.ts +++ b/apps/server/src/dictation/modelFiles.ts @@ -73,6 +73,23 @@ export function rankModels( .map((entry) => entry.path); } +const installedSpeechModels = Effect.fn("dictation.installedSpeechModels")(function* ( + directories: ReadonlyArray, +) { + const fs = yield* FileSystem.FileSystem; + const path = yield* Path.Path; + const found: Array = []; + for (const directory of directories) { + const names = yield* fs.readDirectory(directory).pipe(Effect.orElseSucceed(() => [])); + for (const name of names) { + if (name.toLowerCase().endsWith(".gguf") && !VAD_MODEL.test(name)) { + found.push(path.join(directory, name)); + } + } + } + return found; +}); + /** The model file to load: `configured` when set, else the best one Hex or Scrawl installed. */ export const findModel = Effect.fn("dictation.findModel")(function* (input: { readonly language: string | undefined; @@ -90,15 +107,7 @@ export const findModel = Effect.fn("dictation.findModel")(function* (input: { } return input.configured; } - const found: Array = []; - for (const directory of input.directories) { - const names = yield* fs.readDirectory(directory).pipe(Effect.orElseSucceed(() => [])); - for (const name of names) { - if (name.toLowerCase().endsWith(".gguf") && !VAD_MODEL.test(name)) { - found.push(path.join(directory, name)); - } - } - } + const found = yield* installedSpeechModels(input.directories); const [best] = rankModels(found, input.language, (file) => path.basename(file)); if (best === undefined) { return yield* new DictationError({ @@ -112,6 +121,17 @@ export const findModel = Effect.fn("dictation.findModel")(function* (input: { return best; }); +/** The Parakeet model for `language`: it times words, which captures need and Canary cannot. */ +export const findTimingModel = Effect.fn("dictation.findTimingModel")(function* (input: { + readonly language: string | undefined; + readonly directories: ReadonlyArray; +}) { + const path = yield* Path.Path; + const found = yield* installedSpeechModels(input.directories); + const parakeet = found.filter((file) => path.basename(file).toLowerCase().startsWith("parakeet")); + return rankModels(parakeet, input.language, (file) => path.basename(file))[0]; +}); + /** The Silero model in the first folder that has one. */ export const findVadModel = Effect.fn("dictation.findVadModel")(function* ( directories: ReadonlyArray,