From 230a17b9a144b0914b3f1313b010e0595950fbef Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Sat, 14 Feb 2026 19:41:00 +0000 Subject: [PATCH 1/8] feat(desktop): add minimal local Parakeet support --- apps/desktop/package.json | 2 +- apps/desktop/public/icons/models/nvidia.svg | 3 + apps/desktop/src/constants/models.ts | 66 +++ apps/desktop/src/db/models.ts | 25 +- .../src/main/managers/recording-manager.ts | 97 +++-- .../src/main/managers/service-manager.ts | 50 ++- .../src/pipeline/core/pipeline-types.ts | 1 + .../transcription/parakeet-provider.ts | 329 +++++++++++++++ .../utils/parakeet-feature-extractor.ts | 333 +++++++++++++++ apps/desktop/src/services/model-service.ts | 396 ++++++++++-------- .../src/services/transcription-service.ts | 68 ++- 11 files changed, 1122 insertions(+), 248 deletions(-) create mode 100644 apps/desktop/public/icons/models/nvidia.svg create mode 100644 apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts create mode 100644 apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts diff --git a/apps/desktop/package.json b/apps/desktop/package.json index 0adea6f2..7717504d 100644 --- a/apps/desktop/package.json +++ b/apps/desktop/package.json @@ -44,7 +44,7 @@ "build:types": "pnpm --filter @amical/types build", "build:swift-helper": "pnpm --filter @amical/swift-helper build", "build:windows-helper": "pnpm --filter @amical/windows-helper build", - "build:native-helper": "node -p \"process.platform === 'darwin' ? 'build:swift-helper' : process.platform === 'win32' ? 'build:windows-helper' : 'echo No native helpers'\" | xargs pnpm run", + "build:native-helper": "node -e \"const { execSync } = require('node:child_process'); const cmd = process.platform === 'darwin' ? 'pnpm run build:swift-helper' : process.platform === 'win32' ? 'pnpm run build:windows-helper' : 'echo No native helpers'; execSync(cmd, { stdio: 'inherit' });\"", "dev": "pnpm start", "download-node": "tsx scripts/download-node-binaries.ts", "download-node:all": "tsx scripts/download-node-binaries.ts --all" diff --git a/apps/desktop/public/icons/models/nvidia.svg b/apps/desktop/public/icons/models/nvidia.svg new file mode 100644 index 00000000..9d467c50 --- /dev/null +++ b/apps/desktop/public/icons/models/nvidia.svg @@ -0,0 +1,3 @@ + + + diff --git a/apps/desktop/src/constants/models.ts b/apps/desktop/src/constants/models.ts index 230867e0..f862319c 100644 --- a/apps/desktop/src/constants/models.ts +++ b/apps/desktop/src/constants/models.ts @@ -7,6 +7,12 @@ export interface AvailableWhisperModel { description: string; downloadUrl: string; filename: string; // Expected filename after download + artifacts?: { + filename: string; + downloadUrl: string; + checksum?: string; + size?: number; + }[]; checksum?: string; // Optional checksum for validation features: { icon: string; @@ -15,9 +21,11 @@ export interface AvailableWhisperModel { speed: number; accuracy: number; setup: "offline" | "cloud"; + runtime: "whisper-local" | "parakeet-onnx" | "cloud"; provider: string; providerIcon: string; modelSize: string; + sourceUrl?: string; } // DownloadedModel type is now imported from the database schema @@ -145,9 +153,61 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 4.5, accuracy: 4.5, setup: "cloud", + runtime: "cloud", provider: "Amical Cloud", providerIcon: "/assets/logo.svg", }, + { + id: "parakeet-ctc-0.6b-int8", + name: "NVIDIA Parakeet 0.6B (Local)", + type: "whisper", + description: + "Local CTC speech model optimized for fast English transcription with ONNX Runtime.", + checksum: + "3cfe22e14a7adf70b7b4ab33109a6fad4d7ca61821ca1f37168dc9b3d04b963b", + filename: "model.int8.onnx", + artifacts: [ + { + filename: "model.int8.onnx", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/model.int8.onnx", + checksum: + "3cfe22e14a7adf70b7b4ab33109a6fad4d7ca61821ca1f37168dc9b3d04b963b", + size: 653436437, + }, + { + filename: "vocab.txt", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/vocab.txt", + }, + ], + downloadUrl: + "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/model.int8.onnx", + size: 653436437, + sizeFormatted: "~623 MB", + modelSize: "~623 MB", + features: [ + { + icon: "bolt", + tooltip: "Fast local transcription", + }, + { + icon: "gauge", + tooltip: "DirectML/CPU ONNX runtime", + }, + { + icon: "languages", + tooltip: "English-first CTC vocabulary", + }, + ], + speed: 4.6, + accuracy: 4.2, + setup: "offline", + runtime: "parakeet-onnx", + provider: "NVIDIA", + providerIcon: "/icons/models/nvidia.svg", + sourceUrl: "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx", + }, { id: "whisper-tiny", name: "Whisper Tiny", @@ -177,6 +237,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 5.0, accuracy: 2.5, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, @@ -209,6 +270,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 4.0, accuracy: 3.0, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, @@ -242,6 +304,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 3.0, accuracy: 3.8, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, @@ -274,6 +337,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 2.0, accuracy: 4.3, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, @@ -306,6 +370,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 1.5, accuracy: 4.7, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, @@ -338,6 +403,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 3.5, accuracy: 4.2, setup: "offline", + runtime: "whisper-local", provider: "OpenAI", providerIcon: "/icons/models/openai_dark.svg", }, diff --git a/apps/desktop/src/db/models.ts b/apps/desktop/src/db/models.ts index fb8ca11b..d8c1f627 100644 --- a/apps/desktop/src/db/models.ts +++ b/apps/desktop/src/db/models.ts @@ -158,8 +158,8 @@ export async function getModelsByIds( } /** - * Sync Local Whisper models with filesystem - * Scans the models directory and syncs database records with actual files + * Sync local speech models with filesystem + * Scans expected model paths and syncs database records with actual files */ export async function syncLocalWhisperModels( modelsDirectory: string, @@ -185,24 +185,19 @@ export async function syncLocalWhisperModels( const existingModels = await getModelsByProvider("local-whisper"); const existingModelMap = new Map(existingModels.map((m) => [m.id, m])); - // Scan the models directory for .bin files - const modelFiles = new Set(); - if (fs.existsSync(modelsDirectory)) { - const files = fs.readdirSync(modelsDirectory); - for (const file of files) { - if (file.endsWith(".bin")) { - modelFiles.add(file); - } - } - } - // Map available models by ID for easy lookup // (we already have them indexed by ID, so we don't need this map) // Process each available model for (const model of availableModels) { - const filePath = path.join(modelsDirectory, model.filename); - const fileExists = modelFiles.has(model.filename); + const candidatePaths = [ + path.join(modelsDirectory, model.filename), + path.join(modelsDirectory, model.id, model.filename), + ]; + const filePath = + candidatePaths.find((candidatePath) => fs.existsSync(candidatePath)) || + candidatePaths[1]; + const fileExists = fs.existsSync(filePath); const existingRecord = existingModelMap.get(model.id); if (fileExists) { diff --git a/apps/desktop/src/main/managers/recording-manager.ts b/apps/desktop/src/main/managers/recording-manager.ts index e902096f..75d06ff1 100644 --- a/apps/desktop/src/main/managers/recording-manager.ts +++ b/apps/desktop/src/main/managers/recording-manager.ts @@ -3,6 +3,7 @@ import { EventEmitter } from "node:events"; import { Mutex } from "async-mutex"; import { logger, logPerformance } from "../logger"; import type { ServiceManager } from "@/main/managers/service-manager"; +import type { TranscriptionService } from "../../services/transcription-service"; import type { RecordingState } from "../../types/recording"; import type { ShortcutManager } from "./shortcut-manager"; import { StreamingWavWriter } from "../../utils/streaming-wav-writer"; @@ -67,6 +68,7 @@ export class RecordingManager extends EventEmitter { // System audio state tracking private systemAudioMuted: boolean = false; + private transcriptionServiceUnavailableNotified: boolean = false; constructor(private serviceManager: ServiceManager) { super(); @@ -256,6 +258,7 @@ export class RecordingManager extends EventEmitter { this.setMode(mode); this.terminationCode = null; this.firstChunkReceived = false; + this.transcriptionServiceUnavailableNotified = false; this.recordingStartedAt = performance.now(); this.recordingStoppedAt = null; this.audioChunks = []; @@ -286,10 +289,10 @@ export class RecordingManager extends EventEmitter { try { // Reset VAD state for fresh speech detection (mutex-protected to avoid // interleaving with retry VAD computation) - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); - await transcriptionService.resetVadForNewSession(); + const transcriptionService = this.getTranscriptionService(); + if (transcriptionService) { + await transcriptionService.resetVadForNewSession(); + } // Refresh accessibility context (TextMarker API for Electron support) // Fire and forget - context will be ready by the time first audio chunk arrives @@ -374,10 +377,10 @@ export class RecordingManager extends EventEmitter { // Cancel streaming for cancel codes (not null, not dismissed) if (code && code !== "dismissed" && sessionId) { try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); - await transcriptionService.cancelStreamingSession(sessionId); + const transcriptionService = this.getTranscriptionService(); + if (transcriptionService) { + await transcriptionService.cancelStreamingSession(sessionId); + } } catch (error) { logger.audio.warn("Failed to cancel streaming session", { error }); } @@ -434,14 +437,14 @@ export class RecordingManager extends EventEmitter { // Also send to transcription if we have a session and not terminated if (this.currentSessionId && !this.terminationCode) { try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); - await transcriptionService.processStreamingChunk({ - sessionId: this.currentSessionId, - audioChunk: chunk, - recordingStartedAt: this.recordingStartedAt || undefined, - }); + const transcriptionService = this.getTranscriptionService(); + if (transcriptionService) { + await transcriptionService.processStreamingChunk({ + sessionId: this.currentSessionId, + audioChunk: chunk, + recordingStartedAt: this.recordingStartedAt || undefined, + }); + } } catch (error) { logger.audio.error("Error processing final chunk:", error); } @@ -466,10 +469,13 @@ export class RecordingManager extends EventEmitter { // Stream to transcription (skip if terminated) if (!this.terminationCode) { + const transcriptionService = this.getTranscriptionService(); + if (!transcriptionService) { + await this.endRecording("error"); + return; + } + try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); await transcriptionService.processStreamingChunk({ sessionId, audioChunk: chunk, @@ -545,10 +551,10 @@ export class RecordingManager extends EventEmitter { if (code === "dismissed") { // Cancel streaming session to prevent memory leak and audio bleed try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); - await transcriptionService.cancelStreamingSession(sessionId); + const transcriptionService = this.getTranscriptionService(); + if (transcriptionService) { + await transcriptionService.cancelStreamingSession(sessionId); + } } catch (error) { logger.audio.warn("Failed to cancel streaming session", { error }); } @@ -565,9 +571,14 @@ export class RecordingManager extends EventEmitter { // NORMAL - get transcription and paste let result = ""; try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); + const transcriptionService = this.getTranscriptionService(); + if (!transcriptionService) { + throw new AppError( + "Transcription service unavailable", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + result = await transcriptionService.finalizeSession({ sessionId, audioFilePath: audioFilePath || undefined, @@ -700,12 +711,12 @@ export class RecordingManager extends EventEmitter { // Cancel streaming session if one exists to prevent memory leak and audio bleed if (this.currentSessionId) { try { - const transcriptionService = this.serviceManager.getService( - "transcriptionService", - ); - await transcriptionService.cancelStreamingSession( - this.currentSessionId, - ); + const transcriptionService = this.getTranscriptionService(); + if (transcriptionService) { + await transcriptionService.cancelStreamingSession( + this.currentSessionId, + ); + } } catch (error) { logger.audio.warn("Failed to cancel streaming session", { error }); } @@ -742,9 +753,31 @@ export class RecordingManager extends EventEmitter { this.audioChunks = []; this.terminationCode = null; this.systemAudioMuted = false; + this.transcriptionServiceUnavailableNotified = false; this.clearTimers(); } + private getTranscriptionService(): TranscriptionService | null { + try { + const transcriptionService = this.serviceManager.getService( + "transcriptionService", + ); + return transcriptionService || null; + } catch (error) { + if (!this.transcriptionServiceUnavailableNotified) { + logger.audio.warn("Transcription service unavailable", { + error: error instanceof Error ? error.message : String(error), + }); + this.emit("widget-notification", { + type: "transcription_failed", + errorCode: ErrorCodes.WORKER_INITIALIZATION_FAILED, + }); + this.transcriptionServiceUnavailableNotified = true; + } + return null; + } + } + /** * Create audio file for recording session */ diff --git a/apps/desktop/src/main/managers/service-manager.ts b/apps/desktop/src/main/managers/service-manager.ts index 5a92f6ea..a5699eb5 100644 --- a/apps/desktop/src/main/managers/service-manager.ts +++ b/apps/desktop/src/main/managers/service-manager.ts @@ -176,7 +176,18 @@ export class ServiceManager { this.nativeBridge, this.onboardingService, ); - await this.transcriptionService.initialize(); + try { + await this.transcriptionService.initialize(); + } catch (error) { + this.telemetryService?.captureException(error, { + source: "service_manager", + stage: "initialize_ai_services_preload", + }); + logger.transcription.error( + "Transcription service preload failed, continuing with lazy initialization", + error, + ); + } logger.transcription.info("Transcription Service initialized", { client: "Pipeline with Whisper", @@ -246,24 +257,29 @@ export class ServiceManager { ); } - const services: ServiceMap = { - posthogClient: this.posthogClient!, - telemetryService: this.telemetryService!, - featureFlagService: this.featureFlagService!, - modelService: this.modelService!, - transcriptionService: this.transcriptionService!, - settingsService: this.settingsService!, - authService: this.authService!, - vadService: this.vadService!, - nativeBridge: this.nativeBridge!, - autoUpdaterService: this.autoUpdaterService!, - recordingManager: this.recordingManager!, - shortcutManager: this.shortcutManager!, - windowManager: this.windowManager!, - onboardingService: this.onboardingService!, + const services = { + posthogClient: this.posthogClient, + telemetryService: this.telemetryService, + featureFlagService: this.featureFlagService, + modelService: this.modelService, + transcriptionService: this.transcriptionService, + settingsService: this.settingsService, + authService: this.authService, + vadService: this.vadService, + nativeBridge: this.nativeBridge, + autoUpdaterService: this.autoUpdaterService, + recordingManager: this.recordingManager, + shortcutManager: this.shortcutManager, + windowManager: this.windowManager, + onboardingService: this.onboardingService, }; - return services[serviceName]; + const service = services[serviceName]; + if (!service) { + throw new Error(`Service '${serviceName}' is not available`); + } + + return service as ServiceMap[K]; } async cleanup(): Promise { diff --git a/apps/desktop/src/pipeline/core/pipeline-types.ts b/apps/desktop/src/pipeline/core/pipeline-types.ts index ed6bb384..f2917d9e 100644 --- a/apps/desktop/src/pipeline/core/pipeline-types.ts +++ b/apps/desktop/src/pipeline/core/pipeline-types.ts @@ -10,6 +10,7 @@ export { PipelineContext, SharedPipelineData } from "./context"; // Context for transcription operations (shared between transcribe and flush) export interface TranscribeContext { sessionId?: string; + modelId?: string; vocabulary?: string[]; accessibilityContext?: GetAccessibilityContextResult | null; previousChunk?: string; diff --git a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts new file mode 100644 index 00000000..8669ea62 --- /dev/null +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -0,0 +1,329 @@ +import * as ort from "onnxruntime-node"; +import * as path from "node:path"; +import { + TranscriptionProvider, + TranscribeParams, + TranscribeContext, +} from "../../core/pipeline-types"; +import { ModelService } from "../../../services/model-service"; +import { logger } from "../../../main/logger"; +import { AppError, ErrorCodes } from "../../../types/error"; +import { extractSpeechFromVad } from "../../utils/vad-audio-filter"; +import { + ParakeetFeatureExtractor, + decodeParakeetCtc, + loadParakeetVocabulary, + ParakeetVocabulary, +} from "../../utils/parakeet-feature-extractor"; + +export class ParakeetProvider implements TranscriptionProvider { + readonly name = "parakeet-local"; + + private session: ort.InferenceSession | null = null; + private outputName: string | null = null; + private currentModelId: string | null = null; + private vocabulary: ParakeetVocabulary | null = null; + + private frameBuffer: Float32Array[] = []; + private frameBufferSpeechProbabilities: number[] = []; + private currentSilenceFrameCount = 0; + + private readonly featureExtractor = new ParakeetFeatureExtractor(); + + private readonly FRAME_SIZE = 512; + private readonly MIN_AUDIO_DURATION_MS = 500; + private readonly MAX_SILENCE_DURATION_MS = 3000; + private readonly SAMPLE_RATE = 16000; + private readonly SPEECH_PROBABILITY_THRESHOLD = 0.2; + + constructor(private readonly modelService: ModelService) {} + + async preloadModel(modelId?: string): Promise { + await this.initializeModel(modelId); + } + + async transcribe(params: TranscribeParams): Promise { + const { audioData, speechProbability = 1, context } = params; + await this.initializeModel(context.modelId); + + this.frameBuffer.push(audioData); + this.frameBufferSpeechProbabilities.push(speechProbability); + + const isSpeech = speechProbability > this.SPEECH_PROBABILITY_THRESHOLD; + if (isSpeech) { + this.currentSilenceFrameCount = 0; + } else { + this.currentSilenceFrameCount++; + } + + if (!this.shouldTranscribe()) { + return ""; + } + + return this.doTranscription(context); + } + + async flush(context: TranscribeContext): Promise { + if (this.frameBuffer.length === 0) { + return ""; + } + + await this.initializeModel(context.modelId); + return this.doTranscription(context); + } + + reset(): void { + this.frameBuffer = []; + this.frameBufferSpeechProbabilities = []; + this.currentSilenceFrameCount = 0; + } + + async dispose(): Promise { + if (this.session) { + await this.session.release(); + this.session = null; + } + this.outputName = null; + this.currentModelId = null; + this.vocabulary = null; + this.reset(); + } + + private async doTranscription(context: TranscribeContext): Promise { + try { + if (!this.session || !this.vocabulary || !this.outputName) { + throw new AppError( + "Parakeet model is not initialized", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const vadProbs = [...this.frameBufferSpeechProbabilities]; + const rawAudio = this.aggregateFrames(); + this.reset(); + + const { audio: speechAudio, segments } = extractSpeechFromVad( + rawAudio, + vadProbs, + ); + + if (speechAudio.length === 0) { + logger.transcription.debug( + "Skipping Parakeet transcription - no speech detected by VAD filter", + ); + return ""; + } + + logger.transcription.debug("Parakeet VAD filtered audio", { + before: rawAudio.length, + after: speechAudio.length, + segments: segments.length, + }); + + const features = this.featureExtractor.extract(speechAudio); + const inputTensor = new ort.Tensor( + "float32", + features.inputFeatures, + features.inputShape, + ); + const lengthTensor = new ort.Tensor( + "int64", + BigInt64Array.from([BigInt(features.featuresLength)]), + [1], + ); + + const results = await this.session.run({ + audio_signal: inputTensor, + length: lengthTensor, + }); + + const logitsTensor = results[this.outputName] as ort.Tensor; + const logits = + logitsTensor.data instanceof Float32Array + ? logitsTensor.data + : Float32Array.from(logitsTensor.data as ArrayLike); + + const text = decodeParakeetCtc( + logits, + logitsTensor.dims, + this.vocabulary.tokens, + this.vocabulary.blankTokenId, + features.featuresLength, + ); + + logger.transcription.debug("Parakeet transcription completed", { + textLength: text.length, + featuresLength: features.featuresLength, + }); + + return text; + } catch (error) { + logger.transcription.error("Parakeet transcription failed", { error }); + if (error instanceof AppError) { + throw error; + } + throw new AppError( + `Parakeet transcription failed: ${error instanceof Error ? error.message : String(error)}`, + ErrorCodes.LOCAL_TRANSCRIPTION_FAILED, + ); + } + } + + private async initializeModel(modelId?: string): Promise { + const requestedId = await this.resolveSelectedParakeetModelId(modelId); + if ( + this.session && + this.vocabulary && + this.outputName && + this.currentModelId === requestedId + ) { + return; + } + + const { modelPath, vocabPath } = await this.resolveModelPaths(requestedId); + + if (this.session) { + await this.session.release(); + this.session = null; + } + + const preferredProviders = + process.platform === "win32" + ? (["dml", "cpu"] as const) + : process.platform === "darwin" + ? (["coreml", "cpu"] as const) + : (["cpu"] as const); + + let session: ort.InferenceSession | null = null; + let providersUsed: readonly string[] = preferredProviders; + + try { + session = await ort.InferenceSession.create(modelPath, { + executionProviders: [...preferredProviders], + }); + } catch (error) { + if (preferredProviders.length > 1) { + logger.transcription.warn( + "Parakeet preferred execution provider unavailable, falling back to CPU", + { + requestedProviders: preferredProviders, + error: error instanceof Error ? error.message : String(error), + }, + ); + session = await ort.InferenceSession.create(modelPath, { + executionProviders: ["cpu"], + }); + providersUsed = ["cpu"]; + } else { + throw error; + } + } + + this.session = session; + this.outputName = + session.outputNames.find((name) => /logprob|logits/i.test(name)) || + session.outputNames[0] || + null; + this.vocabulary = await loadParakeetVocabulary(vocabPath); + this.currentModelId = requestedId; + + logger.transcription.info("Initialized local Parakeet model", { + modelId: requestedId, + modelPath, + vocabPath, + executionProviders: providersUsed, + outputName: this.outputName, + }); + } + + private async resolveSelectedParakeetModelId( + requestedModelId?: string, + ): Promise { + const selectedId = requestedModelId || (await this.modelService.getSelectedModel()); + if (!selectedId) { + throw new AppError( + "No speech model selected", + ErrorCodes.MODEL_MISSING, + ); + } + + if (!selectedId.startsWith("parakeet-")) { + throw new AppError( + `Selected model is not a local Parakeet model: ${selectedId}`, + ErrorCodes.MODEL_MISSING, + ); + } + + return selectedId; + } + + private async resolveModelPaths(modelId: string): Promise<{ + modelPath: string; + vocabPath: string; + }> { + const downloadedModels = await this.modelService.getDownloadedModels(); + const downloaded = downloadedModels[modelId]; + + if (!downloaded?.localPath) { + throw new AppError( + `Parakeet model not downloaded: ${modelId}`, + ErrorCodes.MODEL_MISSING, + ); + } + + const modelPath = downloaded.localPath; + const modelDir = path.dirname(modelPath); + + const localFiles = + downloaded.originalModel && + typeof downloaded.originalModel === "object" && + Array.isArray((downloaded.originalModel as { localFiles?: unknown }).localFiles) + ? ((downloaded.originalModel as { localFiles: unknown[] }).localFiles.filter( + (value): value is string => typeof value === "string", + ) as string[]) + : []; + + const vocabPath = + localFiles.find((filePath) => filePath.endsWith("vocab.txt")) || + path.join(modelDir, "vocab.txt"); + + return { modelPath, vocabPath }; + } + + private shouldTranscribe(): boolean { + const audioDurationMs = + ((this.frameBuffer.length * this.FRAME_SIZE) / this.SAMPLE_RATE) * 1000; + const silenceDurationMs = + ((this.currentSilenceFrameCount * this.FRAME_SIZE) / this.SAMPLE_RATE) * + 1000; + + if ( + audioDurationMs >= this.MIN_AUDIO_DURATION_MS && + silenceDurationMs > this.MAX_SILENCE_DURATION_MS + ) { + return true; + } + + if (audioDurationMs > 30000) { + return true; + } + + return false; + } + + private aggregateFrames(): Float32Array { + const totalLength = this.frameBuffer.reduce( + (sum, frame) => sum + frame.length, + 0, + ); + const aggregated = new Float32Array(totalLength); + + let offset = 0; + for (const frame of this.frameBuffer) { + aggregated.set(frame, offset); + offset += frame.length; + } + + return aggregated; + } +} diff --git a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts new file mode 100644 index 00000000..8d8115cf --- /dev/null +++ b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts @@ -0,0 +1,333 @@ +import { promises as fs } from "node:fs"; + +const SAMPLE_RATE = 16000; +const N_FFT = 512; +const WIN_LENGTH = 400; +const HOP_LENGTH = 160; +const PREEMPHASIS = 0.97; +const LOG_ZERO_GUARD = Math.pow(2, -24); +const N_MELS = 80; +const F_MIN = 0; +const F_MAX = SAMPLE_RATE / 2; +const DECODE_SPACE_PATTERN = /^\s|\s\B|(\s)\b/g; + +export interface ParakeetFeatures { + inputFeatures: Float32Array; + inputShape: [number, number, number]; + featuresLength: number; +} + +export interface ParakeetVocabulary { + tokens: string[]; + blankTokenId: number; +} + +function hzToMel(hz: number): number { + return 2595 * Math.log10(1 + hz / 700); +} + +function melToHz(mel: number): number { + return 700 * (Math.pow(10, mel / 2595) - 1); +} + +function buildCenteredHannWindow(): Float32Array { + const window = new Float32Array(N_FFT); + const pad = (N_FFT - WIN_LENGTH) / 2; + for (let i = 0; i < WIN_LENGTH; i++) { + window[pad + i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (WIN_LENGTH - 1)); + } + return window; +} + +function buildMelFilterBank(): Float32Array[] { + const numBins = Math.floor(N_FFT / 2) + 1; + const fbanks: Float32Array[] = Array.from( + { length: N_MELS }, + () => new Float32Array(numBins), + ); + + const minMel = hzToMel(F_MIN); + const maxMel = hzToMel(F_MAX); + const melPoints = new Float64Array(N_MELS + 2); + for (let i = 0; i < melPoints.length; i++) { + melPoints[i] = minMel + ((maxMel - minMel) * i) / (N_MELS + 1); + } + + const bins = new Int32Array(N_MELS + 2); + for (let i = 0; i < bins.length; i++) { + const hz = melToHz(melPoints[i]); + bins[i] = Math.floor(((N_FFT + 1) * hz) / SAMPLE_RATE); + } + + for (let m = 1; m <= N_MELS; m++) { + const left = bins[m - 1]; + const center = bins[m]; + const right = bins[m + 1]; + + if (center > left) { + for (let k = left; k < center && k < numBins; k++) { + fbanks[m - 1][k] = (k - left) / (center - left); + } + } + + if (right > center) { + for (let k = center; k < right && k < numBins; k++) { + fbanks[m - 1][k] = (right - k) / (right - center); + } + } + } + + return fbanks; +} + +function createBitReverseTable(size: number): Uint16Array { + const bits = Math.log2(size); + const table = new Uint16Array(size); + + for (let i = 0; i < size; i++) { + let value = i; + let reversed = 0; + for (let b = 0; b < bits; b++) { + reversed = (reversed << 1) | (value & 1); + value >>= 1; + } + table[i] = reversed; + } + + return table; +} + +function createTwiddleTables(size: number): { + cos: Float32Array; + sin: Float32Array; +} { + const half = size / 2; + const cos = new Float32Array(half); + const sin = new Float32Array(half); + + for (let i = 0; i < half; i++) { + const angle = (2 * Math.PI * i) / size; + cos[i] = Math.cos(angle); + sin[i] = Math.sin(angle); + } + + return { cos, sin }; +} + +function fftInPlace( + real: Float32Array, + imag: Float32Array, + bitReverse: Uint16Array, + cos: Float32Array, + sin: Float32Array, +): void { + const n = real.length; + + for (let i = 0; i < n; i++) { + const j = bitReverse[i]; + if (j > i) { + const tmpR = real[i]; + real[i] = real[j]; + real[j] = tmpR; + + const tmpI = imag[i]; + imag[i] = imag[j]; + imag[j] = tmpI; + } + } + + for (let size = 2; size <= n; size <<= 1) { + const half = size >> 1; + const step = n / size; + + for (let start = 0; start < n; start += size) { + for (let offset = 0; offset < half; offset++) { + const even = start + offset; + const odd = even + half; + const tw = offset * step; + + const wr = cos[tw]; + const wi = sin[tw]; + + const oddR = real[odd]; + const oddI = imag[odd]; + + const tR = oddR * wr + oddI * wi; + const tI = oddI * wr - oddR * wi; + + real[odd] = real[even] - tR; + imag[odd] = imag[even] - tI; + real[even] += tR; + imag[even] += tI; + } + } + } +} + +export class ParakeetFeatureExtractor { + private readonly window = buildCenteredHannWindow(); + private readonly melBanks = buildMelFilterBank(); + private readonly bitReverse = createBitReverseTable(N_FFT); + private readonly twiddle = createTwiddleTables(N_FFT); + + extract(audioData: Float32Array): ParakeetFeatures { + const preemphasized = new Float32Array(audioData.length); + if (audioData.length > 0) { + preemphasized[0] = audioData[0]; + for (let i = 1; i < audioData.length; i++) { + preemphasized[i] = audioData[i] - PREEMPHASIS * audioData[i - 1]; + } + } + + const padded = new Float32Array(preemphasized.length + N_FFT); + padded.set(preemphasized, N_FFT / 2); + + const frameCount = Math.max(1, Math.floor((padded.length - N_FFT) / HOP_LENGTH) + 1); + const featuresLength = Math.max(1, Math.floor(audioData.length / HOP_LENGTH)); + + const logMel = new Float32Array(frameCount * N_MELS); + const real = new Float32Array(N_FFT); + const imag = new Float32Array(N_FFT); + + for (let frame = 0; frame < frameCount; frame++) { + const start = frame * HOP_LENGTH; + real.fill(0); + imag.fill(0); + + for (let i = 0; i < N_FFT; i++) { + real[i] = padded[start + i] * this.window[i]; + } + + fftInPlace( + real, + imag, + this.bitReverse, + this.twiddle.cos, + this.twiddle.sin, + ); + + for (let m = 0; m < N_MELS; m++) { + const bank = this.melBanks[m]; + let energy = 0; + for (let k = 0; k < bank.length; k++) { + const power = real[k] * real[k] + imag[k] * imag[k]; + energy += power * bank[k]; + } + logMel[frame * N_MELS + m] = Math.log(energy + LOG_ZERO_GUARD); + } + } + + const validFrames = Math.min(featuresLength, frameCount); + const normalized = new Float32Array(N_MELS * frameCount); + + for (let m = 0; m < N_MELS; m++) { + let mean = 0; + for (let f = 0; f < validFrames; f++) { + mean += logMel[f * N_MELS + m]; + } + mean /= validFrames; + + let variance = 0; + for (let f = 0; f < validFrames; f++) { + const delta = logMel[f * N_MELS + m] - mean; + variance += delta * delta; + } + const denom = Math.max(validFrames - 1, 1); + variance /= denom; + + const invStd = 1 / (Math.sqrt(variance) + 1e-5); + for (let f = 0; f < frameCount; f++) { + normalized[m * frameCount + f] = + f < validFrames ? (logMel[f * N_MELS + m] - mean) * invStd : 0; + } + } + + return { + inputFeatures: normalized, + inputShape: [1, N_MELS, frameCount], + featuresLength: validFrames, + }; + } +} + +export async function loadParakeetVocabulary( + vocabPath: string, +): Promise { + const content = await fs.readFile(vocabPath, "utf8"); + const tokensById = new Map(); + + for (const rawLine of content.split(/\r?\n/)) { + const line = rawLine.trim(); + if (!line) continue; + + const lastSpace = line.lastIndexOf(" "); + if (lastSpace <= 0) continue; + + const token = line.slice(0, lastSpace).replace(/\u2581/g, " "); + const id = Number.parseInt(line.slice(lastSpace + 1), 10); + if (Number.isNaN(id)) continue; + + tokensById.set(id, token); + } + + const maxId = Math.max(...tokensById.keys()); + const tokens: string[] = Array.from({ length: maxId + 1 }, () => ""); + for (const [id, token] of tokensById.entries()) { + tokens[id] = token; + } + + const blankTokenId = tokens.findIndex((token) => token === ""); + + return { + tokens, + blankTokenId: blankTokenId >= 0 ? blankTokenId : tokens.length - 1, + }; +} + +export function decodeParakeetCtc( + logits: Float32Array, + dims: readonly number[], + vocab: string[], + blankTokenId: number, + requestedLength: number, +): string { + if (dims.length !== 3 || dims[0] !== 1) { + return ""; + } + + const [_, dim1, dim2] = dims; + const vocabSize = vocab.length; + + const timeFirst = dim2 === vocabSize; + const timeSteps = timeFirst ? dim1 : dim2; + const availableLength = Math.min(requestedLength, timeSteps); + + const tokenIds: number[] = []; + let prevId = blankTokenId; + + for (let t = 0; t < availableLength; t++) { + let bestId = 0; + let bestLogit = -Number.MAX_VALUE; + + for (let v = 0; v < vocabSize; v++) { + const index = timeFirst + ? t * vocabSize + v + : v * timeSteps + t; + const value = logits[index] ?? -Number.MAX_VALUE; + if (value > bestLogit) { + bestLogit = value; + bestId = v; + } + } + + if (bestId !== blankTokenId && bestId !== prevId) { + tokenIds.push(bestId); + } + prevId = bestId; + } + + const text = tokenIds.map((id) => vocab[id] ?? "").join(""); + return text.replace(DECODE_SPACE_PATTERN, (_match, capturedSpace) => { + return capturedSpace ? " " : ""; + }); +} diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index 06f88440..23812905 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -60,6 +60,15 @@ class ModelService extends EventEmitter { private state: ModelManagerState; private modelsDirectory: string; private settingsService: SettingsService; + private readonly localSpeechPreference = [ + "whisper-large-v3-turbo", + "parakeet-ctc-0.6b-int8", + "whisper-large-v3", + "whisper-medium", + "whisper-small", + "whisper-base", + "whisper-tiny", + ]; constructor(settingsService: SettingsService) { super(); @@ -105,17 +114,19 @@ class ModelService extends EventEmitter { // Initialize and validate models on startup async initialize(): Promise { try { - // Sync Whisper models with filesystem - const whisperModelsData = AVAILABLE_MODELS.map((model) => ({ - id: model.id, - name: model.name, - description: model.description, - size: model.sizeFormatted, - checksum: model.checksum, - speed: model.speed, - accuracy: model.accuracy, - filename: model.filename, - })); + // Sync local speech models with filesystem + const whisperModelsData = AVAILABLE_MODELS + .filter((model) => model.setup === "offline" && !!model.filename) + .map((model) => ({ + id: model.id, + name: model.name, + description: model.description, + size: model.sizeFormatted, + checksum: model.checksum, + speed: model.speed, + accuracy: model.accuracy, + filename: model.filename, + })); const syncResult = await syncLocalWhisperModels( this.modelsDirectory, @@ -148,22 +159,8 @@ class ModelService extends EventEmitter { const downloadedModelIds = Object.keys(downloadedModels); if (downloadedModelIds.length > 0) { - const preferredOrder = [ - "whisper-large-v3-turbo", - "whisper-large-v3", - "whisper-medium", - "whisper-small", - "whisper-base", - "whisper-tiny", - ]; - - let newModelId = downloadedModelIds[0]; - for (const candidateId of preferredOrder) { - if (downloadedModels[candidateId]) { - newModelId = candidateId; - break; - } - } + const newModelId = + this.pickPreferredLocalModelId(downloadedModelIds); await this.applySpeechModelSelection( newModelId, @@ -198,29 +195,18 @@ class ModelService extends EventEmitter { if (downloadedModelCount > 0) { // Auto-select the best available model using the preferred order - const preferredOrder = [ - "whisper-large-v3-turbo", - "whisper-large-v3", - "whisper-medium", - "whisper-small", - "whisper-base", - "whisper-tiny", - ]; - - for (const candidateId of preferredOrder) { - if (downloadedModels[candidateId]) { - await this.applySpeechModelSelection( - candidateId, - "auto-first-download", - null, - ); - logger.main.info("Auto-selected speech model on initialization", { - modelId: candidateId, - availableModels: Object.keys(downloadedModels), - }); - break; - } - } + const downloadedModelIds = Object.keys(downloadedModels); + const candidateId = + this.pickPreferredLocalModelId(downloadedModelIds); + await this.applySpeechModelSelection( + candidateId, + "auto-first-download", + null, + ); + logger.main.info("Auto-selected speech model on initialization", { + modelId: candidateId, + availableModels: downloadedModelIds, + }); } } @@ -257,22 +243,8 @@ class ModelService extends EventEmitter { if (downloadedModelIds.length > 0) { // Find the best local model from preferred order - const preferredOrder = [ - "whisper-large-v3-turbo", - "whisper-large-v3", - "whisper-medium", - "whisper-small", - "whisper-base", - "whisper-tiny", - ]; - - let newModelId = downloadedModelIds[0]; // Fallback to first available - for (const candidateId of preferredOrder) { - if (downloadedModels[candidateId]) { - newModelId = candidateId; - break; - } - } + const newModelId = + this.pickPreferredLocalModelId(downloadedModelIds); await this.applySpeechModelSelection( newModelId, @@ -318,6 +290,15 @@ class ModelService extends EventEmitter { } } + private pickPreferredLocalModelId(downloadedModelIds: string[]): string { + for (const candidateId of this.localSpeechPreference) { + if (downloadedModelIds.includes(candidateId)) { + return candidateId; + } + } + return downloadedModelIds[0]; + } + // Get all available models from manifest getAvailableModels(): AvailableWhisperModel[] { return AVAILABLE_MODELS; @@ -365,6 +346,10 @@ class ModelService extends EventEmitter { throw new Error(`Model not found: ${modelId}`); } + if (model.setup === "cloud") { + throw new Error(`Cloud model cannot be downloaded: ${modelId}`); + } + if (await this.isModelDownloaded(modelId)) { throw new Error(`Model already downloaded: ${modelId}`); } @@ -374,14 +359,37 @@ class ModelService extends EventEmitter { } const abortController = new AbortController(); - const downloadPath = path.join(this.modelsDirectory, model.filename); + const modelDirectory = path.join(this.modelsDirectory, model.id); + fs.mkdirSync(modelDirectory, { recursive: true }); + + const artifacts = + model.artifacts && model.artifacts.length > 0 + ? model.artifacts + : [ + { + filename: model.filename, + downloadUrl: model.downloadUrl, + checksum: model.checksum, + size: model.size, + }, + ]; + const primaryArtifact = + artifacts.find((artifact) => artifact.filename === model.filename) || + artifacts[0]; + const downloadPath = path.join(modelDirectory, primaryArtifact.filename); const progress: DownloadProgress = { modelId, progress: 0, status: "downloading", bytesDownloaded: 0, - totalBytes: model.size, + totalBytes: (() => { + const artifactBytes = artifacts.reduce( + (sum, artifact) => sum + (artifact.size || 0), + 0, + ); + return artifactBytes > 0 ? artifactBytes : model.size; + })(), abortController, }; @@ -392,85 +400,103 @@ class ModelService extends EventEmitter { logger.main.info("Starting model download", { modelId, size: model.sizeFormatted, - url: model.downloadUrl, + artifacts: artifacts.map((artifact) => artifact.filename), }); - const response = await fetch(model.downloadUrl, { - signal: abortController.signal, - headers: { - "User-Agent": getUserAgent(), - }, - }); - - if (!response.ok) { - throw new Error( - `Failed to download: ${response.status} ${response.statusText}`, - ); - } - - const totalBytes = - parseInt(response.headers.get("content-length") || "0") || model.size; - progress.totalBytes = totalBytes; - - const fileStream = fs.createWriteStream(downloadPath); let bytesDownloaded = 0; let lastProgressEmit = 0; + const localFiles: string[] = []; - const reader = response.body?.getReader(); - if (!reader) { - throw new Error("Failed to get response reader"); - } + for (const artifact of artifacts) { + const artifactPath = path.join(modelDirectory, artifact.filename); - while (true) { - const { done, value } = await reader.read(); + const response = await fetch(artifact.downloadUrl, { + signal: abortController.signal, + headers: { + "User-Agent": getUserAgent(), + }, + }); - if (done) break; + if (!response.ok) { + throw new Error( + `Failed to download ${artifact.filename}: ${response.status} ${response.statusText}`, + ); + } + + const artifactBytes = + parseInt(response.headers.get("content-length") || "0") || + artifact.size || + 0; + if (!artifact.size && artifactBytes > 0) { + progress.totalBytes += artifactBytes; + } - if (abortController.signal.aborted) { - fileStream.close(); - fs.unlinkSync(downloadPath); - throw new Error("Download cancelled"); + const fileStream = fs.createWriteStream(artifactPath); + const reader = response.body?.getReader(); + if (!reader) { + throw new Error(`Failed to read ${artifact.filename}`); } - fileStream.write(value); - bytesDownloaded += value.length; + while (true) { + const { done, value } = await reader.read(); + if (done) break; - progress.bytesDownloaded = bytesDownloaded; - progress.progress = Math.round((bytesDownloaded / totalBytes) * 100); + if (abortController.signal.aborted) { + fileStream.close(); + if (fs.existsSync(artifactPath)) { + fs.unlinkSync(artifactPath); + } + throw new Error("Download cancelled"); + } - // Emit progress every 1% or 1MB to avoid too many events - const progressPercent = progress.progress; - if ( - progressPercent - lastProgressEmit >= 1 || - bytesDownloaded - (lastProgressEmit * totalBytes) / 100 >= 1024 * 1024 - ) { - this.emit("download-progress", modelId, { ...progress }); - lastProgressEmit = progressPercent; + fileStream.write(value); + bytesDownloaded += value.length; + progress.bytesDownloaded = bytesDownloaded; + progress.progress = + progress.totalBytes > 0 + ? Math.round((bytesDownloaded / progress.totalBytes) * 100) + : 0; + + const progressPercent = progress.progress; + if ( + progressPercent - lastProgressEmit >= 1 || + bytesDownloaded - (lastProgressEmit * progress.totalBytes) / 100 >= + 1024 * 1024 + ) { + this.emit("download-progress", modelId, { ...progress }); + lastProgressEmit = progressPercent; + } } - } - fileStream.end(); + await new Promise((resolve, reject) => { + fileStream.end(() => resolve()); + fileStream.on("error", reject); + }); + + if (artifact.checksum) { + const fileChecksum = await this.calculateFileChecksum( + artifactPath, + artifact.checksum, + ); + if (fileChecksum !== artifact.checksum.toLowerCase()) { + fs.unlinkSync(artifactPath); + throw new Error( + `Checksum mismatch for ${artifact.filename}. Expected: ${artifact.checksum}, Got: ${fileChecksum}`, + ); + } + } + + localFiles.push(artifactPath); + } - // Get actual file size (no validation against expected size) const stats = fs.statSync(downloadPath); logger.main.info("Download completed", { modelId, - expectedSize: totalBytes, + fileCount: localFiles.length, + primaryPath: downloadPath, actualSize: stats.size, - sizeDifference: Math.abs(stats.size - totalBytes), }); - // Verify checksum if provided - if (model.checksum) { - const fileChecksum = await this.calculateFileChecksum(downloadPath); - if (fileChecksum !== model.checksum) { - fs.unlinkSync(downloadPath); - throw new Error( - `Checksum mismatch. Expected: ${model.checksum}, Got: ${fileChecksum}`, - ); - } - } - // Create/update model record in database with download info await upsertModel({ id: model.id, @@ -483,10 +509,19 @@ class ModelService extends EventEmitter { speed: model.speed, accuracy: model.accuracy, localPath: downloadPath, - sizeBytes: stats.size, + sizeBytes: localFiles.reduce((sum, filePath) => { + try { + return sum + fs.statSync(filePath).size; + } catch { + return sum; + } + }, 0), downloadedAt: new Date(), context: null, - originalModel: null, + originalModel: { + localFiles, + sourceUrl: model.sourceUrl || null, + }, }); // Get the updated model from database @@ -504,7 +539,7 @@ class ModelService extends EventEmitter { logger.main.info("Model download completed", { modelId, path: downloadPath, - size: stats.size, + size: downloadedModel.sizeBytes, }); // Auto-select if this is the first model @@ -526,9 +561,9 @@ class ModelService extends EventEmitter { } catch (error) { // Clean up on error this.state.activeDownloads.delete(modelId); - - if (fs.existsSync(downloadPath)) { - fs.unlinkSync(downloadPath); + const modelDirectory = path.join(this.modelsDirectory, model.id); + if (fs.existsSync(modelDirectory)) { + fs.rmSync(modelDirectory, { recursive: true, force: true }); } const err = error instanceof Error ? error : new Error(String(error)); @@ -579,12 +614,36 @@ class ModelService extends EventEmitter { const wasSelected = currentSelection === modelId; // Delete file - if (downloadedModel.localPath && fs.existsSync(downloadedModel.localPath)) { - fs.unlinkSync(downloadedModel.localPath); - logger.main.info("Deleted model file", { - modelId, - path: downloadedModel.localPath, - }); + const localFiles = + downloadedModel.originalModel && + typeof downloadedModel.originalModel === "object" && + Array.isArray( + (downloadedModel.originalModel as { localFiles?: unknown }).localFiles, + ) + ? ( + downloadedModel.originalModel as { + localFiles: unknown[]; + } + ).localFiles.filter( + (value): value is string => typeof value === "string", + ) + : downloadedModel.localPath + ? [downloadedModel.localPath] + : []; + + for (const localFile of localFiles) { + if (fs.existsSync(localFile)) { + fs.unlinkSync(localFile); + logger.main.info("Deleted model file", { + modelId, + path: localFile, + }); + } + } + + const modelDirectory = path.join(this.modelsDirectory, modelId); + if (fs.existsSync(modelDirectory)) { + fs.rmSync(modelDirectory, { recursive: true, force: true }); } // Remove the model record from database (we only store downloaded models) @@ -594,30 +653,24 @@ class ModelService extends EventEmitter { if (wasSelected) { // Try to auto-select next best model const remainingModels = await this.getValidDownloadedModels(); - const preferredOrder = [ - "whisper-large-v3-turbo", - "whisper-large-v1", - "whisper-medium", - "whisper-small", - "whisper-base", - "whisper-tiny", - ]; + const remainingModelIds = Object.keys(remainingModels); + const candidateId = + remainingModelIds.length > 0 + ? this.pickPreferredLocalModelId(remainingModelIds) + : null; let autoSelected = false; - for (const candidateId of preferredOrder) { - if (remainingModels[candidateId]) { - await this.applySpeechModelSelection( - candidateId, - "auto-after-deletion", - modelId, - ); - logger.main.info("Auto-selected new model after deletion", { - oldModel: modelId, - newModel: candidateId, - }); - autoSelected = true; - break; - } + if (candidateId) { + await this.applySpeechModelSelection( + candidateId, + "auto-after-deletion", + modelId, + ); + logger.main.info("Auto-selected new model after deletion", { + oldModel: modelId, + newModel: candidateId, + }); + autoSelected = true; } if (!autoSelected) { @@ -635,14 +688,19 @@ class ModelService extends EventEmitter { await this.validateAndClearInvalidDefaults(); } - // Calculate file checksum (SHA-1) - private async calculateFileChecksum(filePath: string): Promise { + // Calculate file checksum (auto-detect algorithm from expected hash length) + private async calculateFileChecksum( + filePath: string, + expectedChecksum?: string, + ): Promise { + const algorithm = + expectedChecksum && expectedChecksum.length === 64 ? "sha256" : "sha1"; return new Promise((resolve, reject) => { - const hash = crypto.createHash("sha1"); + const hash = crypto.createHash(algorithm); const stream = fs.createReadStream(filePath); stream.on("data", (data) => hash.update(data)); - stream.on("end", () => resolve(hash.digest("hex"))); + stream.on("end", () => resolve(hash.digest("hex").toLowerCase())); stream.on("error", reject); }); } @@ -681,8 +739,10 @@ class ModelService extends EventEmitter { (await this.settingsService.getFormatterConfig()) || { enabled: false }; const currentModelId = formatterConfig.modelId; const fallbackModelId = formatterConfig.fallbackModelId; - const movedToCloud = newModelId === "amical-cloud"; - const movedFromCloud = oldModelId === "amical-cloud"; + const isCloudSpeechModelId = (modelId: string | null | undefined) => + !!AVAILABLE_MODELS.find((m) => m.id === modelId && m.setup === "cloud"); + const movedToCloud = isCloudSpeechModelId(newModelId); + const movedFromCloud = isCloudSpeechModelId(oldModelId); const usingCloudFormatting = currentModelId === "amical-cloud"; let nextConfig = { ...formatterConfig }; @@ -787,7 +847,7 @@ class ModelService extends EventEmitter { // Otherwise, find the best available model (prioritize by quality) const preferredOrder = [ "whisper-large-v3-turbo", - "whisper-large-v1", + "whisper-large-v3", "whisper-medium", "whisper-small", "whisper-base", @@ -1130,7 +1190,7 @@ class ModelService extends EventEmitter { const availableModel = AVAILABLE_MODELS.find( (m) => m.id === defaultSpeechModel, ); - const isAmicalModel = availableModel?.provider === "Amical Cloud"; + const isAmicalModel = availableModel?.setup === "cloud"; const existsInDb = await modelExists("local-whisper", defaultSpeechModel); // Amical cloud models are always valid; local models must exist in DB diff --git a/apps/desktop/src/services/transcription-service.ts b/apps/desktop/src/services/transcription-service.ts index 1ca6221e..2d2f0659 100644 --- a/apps/desktop/src/services/transcription-service.ts +++ b/apps/desktop/src/services/transcription-service.ts @@ -7,6 +7,7 @@ import { } from "../pipeline/core/pipeline-types"; import { createDefaultContext } from "../pipeline/core/context"; import { WhisperProvider } from "../pipeline/providers/transcription/whisper-provider"; +import { ParakeetProvider } from "../pipeline/providers/transcription/parakeet-provider"; import { AmicalCloudProvider } from "../pipeline/providers/transcription/amical-cloud-provider"; import { OpenRouterProvider } from "../pipeline/providers/formatting/openrouter-formatter"; import { OllamaFormatter } from "../pipeline/providers/formatting/ollama-formatter"; @@ -36,6 +37,7 @@ import * as fs from "node:fs"; */ export class TranscriptionService { private whisperProvider: WhisperProvider; + private parakeetProvider: ParakeetProvider; private cloudProvider: AmicalCloudProvider; private currentProvider: TranscriptionProvider | null = null; private streamingSessions = new Map(); @@ -58,6 +60,7 @@ export class TranscriptionService { private onboardingService: OnboardingService | null, ) { this.whisperProvider = new WhisperProvider(modelService); + this.parakeetProvider = new ParakeetProvider(modelService); this.cloudProvider = new AmicalCloudProvider(); this.vadService = vadService; this.settingsService = settingsService; @@ -83,12 +86,18 @@ export class TranscriptionService { // Find the model in AVAILABLE_MODELS const model = AVAILABLE_MODELS.find((m) => m.id === selectedModelId); - // Use cloud provider for Amical Cloud models - if (model?.provider === "Amical Cloud") { + // Use cloud provider for cloud-backed models + if (model?.setup === "cloud") { this.currentProvider = this.cloudProvider; return this.cloudProvider; } + // Use Parakeet provider for local Parakeet ONNX models + if (model?.runtime === "parakeet-onnx") { + this.currentProvider = this.parakeetProvider; + return this.parakeetProvider; + } + // Default to whisper for all other models this.currentProvider = this.whisperProvider; return this.whisperProvider; @@ -100,7 +109,7 @@ export class TranscriptionService { const model = selectedModelId ? AVAILABLE_MODELS.find((m) => m.id === selectedModelId) : null; - const isCloudModel = model?.provider === "Amical Cloud"; + const isCloudModel = model?.setup === "cloud"; // Only preload for local models if (!isCloudModel) { @@ -114,10 +123,18 @@ export class TranscriptionService { // Check if models are available for preloading const hasModels = await this.isModelAvailable(); if (hasModels) { - logger.transcription.info("Preloading Whisper model..."); - await this.preloadWhisperModel(); - this.modelWasPreloaded = true; - logger.transcription.info("Whisper model preloaded successfully"); + const provider = await this.selectProvider(); + if (provider === this.parakeetProvider) { + logger.transcription.info("Preloading Parakeet model..."); + await this.parakeetProvider.preloadModel(selectedModelId || undefined); + this.modelWasPreloaded = true; + logger.transcription.info("Parakeet model preloaded successfully"); + } else { + logger.transcription.info("Preloading Whisper model..."); + await this.preloadWhisperModel(); + this.modelWasPreloaded = true; + logger.transcription.info("Whisper model preloaded successfully"); + } } else { logger.transcription.info( "Whisper model preloading skipped - no models available", @@ -172,7 +189,7 @@ export class TranscriptionService { const selectedModelId = await this.modelService.getSelectedModel(); if (selectedModelId) { const model = AVAILABLE_MODELS.find((m) => m.id === selectedModelId); - if (model?.provider === "Amical Cloud") { + if (model?.setup === "cloud") { return true; } } @@ -206,12 +223,25 @@ export class TranscriptionService { if (shouldPreload) { const hasModels = await this.isModelAvailable(); if (hasModels) { - logger.transcription.info( - "Loading Whisper model after model change...", - ); - await this.whisperProvider.preloadModel(); - this.modelWasPreloaded = true; - logger.transcription.info("Whisper model loaded successfully"); + const provider = await this.selectProvider(); + const selectedModelId = await this.modelService.getSelectedModel(); + if (provider === this.parakeetProvider) { + logger.transcription.info( + "Loading Parakeet model after model change...", + ); + await this.parakeetProvider.preloadModel( + selectedModelId || undefined, + ); + this.modelWasPreloaded = true; + logger.transcription.info("Parakeet model loaded successfully"); + } else { + logger.transcription.info( + "Loading Whisper model after model change...", + ); + await this.whisperProvider.preloadModel(); + this.modelWasPreloaded = true; + logger.transcription.info("Whisper model loaded successfully"); + } } else { logger.transcription.info("No models available to preload"); } @@ -312,6 +342,7 @@ export class TranscriptionService { // Select the appropriate provider const provider = await this.selectProvider(); + const selectedSpeechModelId = await this.modelService.getSelectedModel(); // Transcribe chunk (flush is done separately in finalizeSession) const chunkTranscription = await provider.transcribe({ @@ -319,6 +350,7 @@ export class TranscriptionService { speechProbability: speechProbability, context: { sessionId, + modelId: selectedSpeechModelId || undefined, vocabulary: session.context.sharedData.vocabulary, accessibilityContext: session.context.sharedData.accessibilityContext, previousChunk, @@ -420,8 +452,10 @@ export class TranscriptionService { const provider = await this.selectProvider(); usedCloudProvider = provider.name === "amical-cloud"; + const selectedSpeechModelId = await this.modelService.getSelectedModel(); const finalTranscription = await provider.flush({ sessionId, + modelId: selectedSpeechModelId || undefined, vocabulary: session.context.sharedData.vocabulary, accessibilityContext: session.context.sharedData.accessibilityContext, previousChunk, @@ -485,12 +519,13 @@ export class TranscriptionService { audioFilePath, hasAudioFile: !!audioFilePath, }); + const selectedSpeechModelId = await this.modelService.getSelectedModel(); await createTranscription({ text: completeTranscription, language: session.context.sharedData.userPreferences?.language || "en", duration: session.context.sharedData.audioMetadata?.duration, - speechModel: "whisper-local", + speechModel: selectedSpeechModelId || "whisper-local", formattingModel, audioFile: audioFilePath, meta: { @@ -925,6 +960,7 @@ export class TranscriptionService { audioData: frames[i], speechProbability: vadProbs[i], context: { + modelId: selectedModelId || undefined, vocabulary, language, previousChunk, @@ -942,6 +978,7 @@ export class TranscriptionService { // Flush to get remaining buffered audio const aggregatedTranscription = transcriptionResults.join(""); const finalTranscription = await provider.flush({ + modelId: selectedModelId || undefined, vocabulary, language, aggregatedTranscription: aggregatedTranscription || undefined, @@ -1028,6 +1065,7 @@ export class TranscriptionService { */ async dispose(): Promise { await this.whisperProvider.dispose(); + await this.parakeetProvider.dispose(); // VAD service is managed by ServiceManager logger.transcription.info("Transcription service disposed"); } From 14e162daea11f6a3cc1fb7e5a67a3fd224b8377d Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Sat, 14 Feb 2026 20:16:33 +0000 Subject: [PATCH 2/8] feat(desktop): add Parakeet TDT support with CTC fallback --- apps/desktop/src/constants/models.ts | 70 +- .../transcription/parakeet-provider.ts | 698 +++++++++++++++--- .../utils/parakeet-feature-extractor.ts | 55 +- apps/desktop/src/services/model-service.ts | 1 + apps/desktop/src/trpc/routers/models.ts | 6 + 5 files changed, 716 insertions(+), 114 deletions(-) diff --git a/apps/desktop/src/constants/models.ts b/apps/desktop/src/constants/models.ts index f862319c..cf1bc351 100644 --- a/apps/desktop/src/constants/models.ts +++ b/apps/desktop/src/constants/models.ts @@ -159,10 +159,10 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ }, { id: "parakeet-ctc-0.6b-int8", - name: "NVIDIA Parakeet 0.6B (Local)", + name: "NVIDIA Parakeet CTC 0.6B", type: "whisper", description: - "Local CTC speech model optimized for fast English transcription with ONNX Runtime.", + "CTC speech model optimized for fast English transcription with ONNX Runtime.", checksum: "3cfe22e14a7adf70b7b4ab33109a6fad4d7ca61821ca1f37168dc9b3d04b963b", filename: "model.int8.onnx", @@ -208,6 +208,72 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ providerIcon: "/icons/models/nvidia.svg", sourceUrl: "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx", }, + { + id: "parakeet-tdt-0.6b-v3-int8", + name: "NVIDIA Parakeet TDT 0.6B v3", + type: "whisper", + description: + "Transducer speech model with improved multilingual quality and robustness using ONNX Runtime.", + checksum: "", + filename: "encoder-model.int8.onnx", + artifacts: [ + { + filename: "encoder-model.int8.onnx", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/encoder-model.int8.onnx", + size: 652183999, + }, + { + filename: "decoder_joint-model.int8.onnx", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/decoder_joint-model.int8.onnx", + size: 18202004, + }, + { + filename: "nemo128.onnx", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/nemo128.onnx", + size: 139764, + }, + { + filename: "vocab.txt", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/vocab.txt", + size: 93939, + }, + { + filename: "config.json", + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/config.json", + }, + ], + downloadUrl: + "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx/resolve/main/encoder-model.int8.onnx", + size: 670619706, + sizeFormatted: "~640 MB", + modelSize: "~640 MB", + features: [ + { + icon: "award", + tooltip: "Higher-quality transducer decoding", + }, + { + icon: "gauge", + tooltip: "DirectML/CPU ONNX runtime", + }, + { + icon: "languages", + tooltip: "Strong multilingual support", + }, + ], + speed: 4.3, + accuracy: 4.6, + setup: "offline", + runtime: "parakeet-onnx", + provider: "NVIDIA", + providerIcon: "/icons/models/nvidia.svg", + sourceUrl: "https://huggingface.co/istupakov/parakeet-tdt-0.6b-v3-onnx", + }, { id: "whisper-tiny", name: "Whisper Tiny", diff --git a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts index 8669ea62..b9b2e7a2 100644 --- a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -1,5 +1,6 @@ import * as ort from "onnxruntime-node"; import * as path from "node:path"; +import { promises as fs } from "node:fs"; import { TranscriptionProvider, TranscribeParams, @@ -12,23 +13,58 @@ import { extractSpeechFromVad } from "../../utils/vad-audio-filter"; import { ParakeetFeatureExtractor, decodeParakeetCtc, + decodeParakeetTokens, loadParakeetVocabulary, ParakeetVocabulary, + ParakeetFeatures, } from "../../utils/parakeet-feature-extractor"; +type ParakeetModelVariant = "ctc" | "tdt"; + +interface ResolvedParakeetPaths { + variant: ParakeetModelVariant; + ctcModelPath?: string; + encoderModelPath?: string; + decoderJointModelPath?: string; + vocabPath: string; + configPath?: string; +} + +interface ParakeetModelConfig { + features_size?: number; + max_tokens_per_step?: number; +} + +interface EncoderAccessor { + hiddenSize: number; + timeSteps: number; + at: (timeStep: number) => Float32Array; +} + +interface TdtDecoderState { + state1: ort.Tensor; + state2: ort.Tensor; +} + export class ParakeetProvider implements TranscriptionProvider { readonly name = "parakeet-local"; - private session: ort.InferenceSession | null = null; - private outputName: string | null = null; + private ctcSession: ort.InferenceSession | null = null; + private tdtEncoderSession: ort.InferenceSession | null = null; + private tdtDecoderJointSession: ort.InferenceSession | null = null; + + private ctcOutputName: string | null = null; private currentModelId: string | null = null; + private currentModelVariant: ParakeetModelVariant | null = null; private vocabulary: ParakeetVocabulary | null = null; private frameBuffer: Float32Array[] = []; private frameBufferSpeechProbabilities: number[] = []; private currentSilenceFrameCount = 0; - private readonly featureExtractor = new ParakeetFeatureExtractor(); + private featureSize = 80; + private maxTokensPerStep = 10; + private featureExtractor = new ParakeetFeatureExtractor(80); private readonly FRAME_SIZE = 512; private readonly MIN_AUDIO_DURATION_MS = 500; @@ -79,19 +115,17 @@ export class ParakeetProvider implements TranscriptionProvider { } async dispose(): Promise { - if (this.session) { - await this.session.release(); - this.session = null; - } - this.outputName = null; + await this.releaseSessions(); + this.ctcOutputName = null; this.currentModelId = null; + this.currentModelVariant = null; this.vocabulary = null; this.reset(); } - private async doTranscription(context: TranscribeContext): Promise { + private async doTranscription(_context: TranscribeContext): Promise { try { - if (!this.session || !this.vocabulary || !this.outputName) { + if (!this.vocabulary || !this.currentModelVariant) { throw new AppError( "Parakeet model is not initialized", ErrorCodes.WORKER_INITIALIZATION_FAILED, @@ -118,45 +152,16 @@ export class ParakeetProvider implements TranscriptionProvider { before: rawAudio.length, after: speechAudio.length, segments: segments.length, + variant: this.currentModelVariant, }); const features = this.featureExtractor.extract(speechAudio); - const inputTensor = new ort.Tensor( - "float32", - features.inputFeatures, - features.inputShape, - ); - const lengthTensor = new ort.Tensor( - "int64", - BigInt64Array.from([BigInt(features.featuresLength)]), - [1], - ); - - const results = await this.session.run({ - audio_signal: inputTensor, - length: lengthTensor, - }); - - const logitsTensor = results[this.outputName] as ort.Tensor; - const logits = - logitsTensor.data instanceof Float32Array - ? logitsTensor.data - : Float32Array.from(logitsTensor.data as ArrayLike); - - const text = decodeParakeetCtc( - logits, - logitsTensor.dims, - this.vocabulary.tokens, - this.vocabulary.blankTokenId, - features.featuresLength, - ); - logger.transcription.debug("Parakeet transcription completed", { - textLength: text.length, - featuresLength: features.featuresLength, - }); + if (this.currentModelVariant === "tdt") { + return this.transcribeTdt(features); + } - return text; + return this.transcribeCtc(features); } catch (error) { logger.transcription.error("Parakeet transcription failed", { error }); if (error instanceof AppError) { @@ -169,23 +174,265 @@ export class ParakeetProvider implements TranscriptionProvider { } } + private async transcribeCtc(features: ParakeetFeatures): Promise { + if (!this.ctcSession || !this.vocabulary || !this.ctcOutputName) { + throw new AppError( + "Parakeet CTC model is not initialized", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const inputName = this.findName(this.ctcSession.inputNames, [/audio_signal/i], 0); + const lengthName = this.findName(this.ctcSession.inputNames, [/length/i], 1); + + const inputTensor = new ort.Tensor( + "float32", + features.inputFeatures, + features.inputShape, + ); + const lengthTensor = this.createIntegerTensorForInput( + this.ctcSession, + lengthName, + [features.featuresLength], + [1], + ); + + const results = await this.ctcSession.run({ + [inputName]: inputTensor, + [lengthName]: lengthTensor, + }); + + const logitsTensor = results[this.ctcOutputName] as ort.Tensor; + const logits = this.toFloat32Array(logitsTensor.data); + + return decodeParakeetCtc( + logits, + logitsTensor.dims, + this.vocabulary.tokens, + this.vocabulary.blankTokenId, + features.featuresLength, + ); + } + + private async transcribeTdt(features: ParakeetFeatures): Promise { + if (!this.tdtEncoderSession || !this.tdtDecoderJointSession || !this.vocabulary) { + throw new AppError( + "Parakeet TDT model is not initialized", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const encoderInputName = this.findName( + this.tdtEncoderSession.inputNames, + [/audio_signal/i], + 0, + ); + const encoderLengthName = this.findName( + this.tdtEncoderSession.inputNames, + [/length/i], + 1, + ); + + const encoderOutName = this.findName( + this.tdtEncoderSession.outputNames, + [/^outputs$/i], + 0, + ); + const encodedLengthName = this.findName( + this.tdtEncoderSession.outputNames, + [/encoded_lengths/i, /length/i], + 1, + ); + + const encoderInputTensor = new ort.Tensor( + "float32", + features.inputFeatures, + features.inputShape, + ); + const encoderLengthTensor = this.createIntegerTensorForInput( + this.tdtEncoderSession, + encoderLengthName, + [features.featuresLength], + [1], + ); + + const encoderResults = await this.tdtEncoderSession.run({ + [encoderInputName]: encoderInputTensor, + [encoderLengthName]: encoderLengthTensor, + }); + + const encoderTensor = encoderResults[encoderOutName] as ort.Tensor; + const encodedLengthTensor = encoderResults[encodedLengthName] as ort.Tensor; + + const accessor = this.createEncoderAccessor(encoderTensor); + const encodedLengths = this.toBigInt64Array(encodedLengthTensor.data); + const encodedLength = Math.max( + 1, + Math.min(accessor.timeSteps, Number(encodedLengths[0] ?? BigInt(accessor.timeSteps))), + ); + + const decoderOutputName = this.findName( + this.tdtDecoderJointSession.outputNames, + [/^outputs$/i], + 0, + ); + const outputState1Name = this.findName( + this.tdtDecoderJointSession.outputNames, + [/output_states_1/i], + 1, + ); + const outputState2Name = this.findName( + this.tdtDecoderJointSession.outputNames, + [/output_states_2/i], + 2, + ); + + const decoderEncoderInputName = this.findName( + this.tdtDecoderJointSession.inputNames, + [/encoder_outputs/i], + 0, + ); + const decoderTargetsInputName = this.findName( + this.tdtDecoderJointSession.inputNames, + [/targets/i], + 1, + ); + const decoderTargetLengthInputName = this.findName( + this.tdtDecoderJointSession.inputNames, + [/target_length/i], + 2, + ); + const inputState1Name = this.findName( + this.tdtDecoderJointSession.inputNames, + [/input_states_1/i], + 3, + ); + const inputState2Name = this.findName( + this.tdtDecoderJointSession.inputNames, + [/input_states_2/i], + 4, + ); + + let state = this.createInitialTdtState( + this.tdtDecoderJointSession, + inputState1Name, + inputState2Name, + accessor.hiddenSize, + ); + + const tokenIds: number[] = []; + const blankTokenId = this.vocabulary.blankTokenId; + + let t = 0; + let emittedTokens = 0; + let guard = 0; + const guardLimit = encodedLength * Math.max(8, this.maxTokensPerStep * 4); + + while (t < encodedLength && guard++ < guardLimit) { + const encoderStepTensor = new ort.Tensor( + "float32", + accessor.at(t), + [1, accessor.hiddenSize, 1], + ); + + const lastTokenId = tokenIds.length > 0 ? tokenIds[tokenIds.length - 1] : blankTokenId; + const targetsTensor = this.createIntegerTensorForInput( + this.tdtDecoderJointSession, + decoderTargetsInputName, + [lastTokenId], + [1, 1], + ); + const targetLengthTensor = this.createIntegerTensorForInput( + this.tdtDecoderJointSession, + decoderTargetLengthInputName, + [1], + [1], + ); + + const decoderResults = await this.tdtDecoderJointSession.run({ + [decoderEncoderInputName]: encoderStepTensor, + [decoderTargetsInputName]: targetsTensor, + [decoderTargetLengthInputName]: targetLengthTensor, + [inputState1Name]: state.state1, + [inputState2Name]: state.state2, + }); + + const outputTensor = decoderResults[decoderOutputName] as ort.Tensor; + const outputData = this.toFloat32Array(outputTensor.data); + const vocabSize = this.vocabulary.tokens.length; + + if (outputData.length < vocabSize) { + throw new AppError( + "Unexpected TDT decoder output shape", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const token = this.argmax(outputData, 0, vocabSize); + const stepCount = + outputData.length > vocabSize + ? this.argmax(outputData, vocabSize, outputData.length - vocabSize) + : 0; + + if (token !== blankTokenId) { + tokenIds.push(token); + emittedTokens++; + + state = { + state1: decoderResults[outputState1Name] as ort.Tensor, + state2: decoderResults[outputState2Name] as ort.Tensor, + }; + } + + if (stepCount > 0) { + t += stepCount; + emittedTokens = 0; + } else if (token === blankTokenId || emittedTokens >= this.maxTokensPerStep) { + t += 1; + emittedTokens = 0; + } + } + + if (guard >= guardLimit) { + logger.transcription.warn("TDT decoding stopped by safety guard", { + encodedLength, + emittedTokens: tokenIds.length, + }); + } + + return decodeParakeetTokens(tokenIds, this.vocabulary.tokens); + } + private async initializeModel(modelId?: string): Promise { const requestedId = await this.resolveSelectedParakeetModelId(modelId); if ( - this.session && this.vocabulary && - this.outputName && - this.currentModelId === requestedId + this.currentModelId === requestedId && + ((this.currentModelVariant === "ctc" && this.ctcSession) || + (this.currentModelVariant === "tdt" && + this.tdtEncoderSession && + this.tdtDecoderJointSession)) ) { return; } - const { modelPath, vocabPath } = await this.resolveModelPaths(requestedId); + const resolved = await this.resolveModelPaths(requestedId); + await this.releaseSessions(); - if (this.session) { - await this.session.release(); - this.session = null; - } + const config = await this.loadModelConfig(resolved.configPath); + this.featureSize = + typeof config.features_size === "number" + ? config.features_size + : resolved.variant === "tdt" + ? 128 + : 80; + this.maxTokensPerStep = + typeof config.max_tokens_per_step === "number" + ? config.max_tokens_per_step + : 10; + this.featureExtractor = new ParakeetFeatureExtractor(this.featureSize); + + this.vocabulary = await loadParakeetVocabulary(resolved.vocabPath); const preferredProviders = process.platform === "win32" @@ -194,46 +441,100 @@ export class ParakeetProvider implements TranscriptionProvider { ? (["coreml", "cpu"] as const) : (["cpu"] as const); - let session: ort.InferenceSession | null = null; - let providersUsed: readonly string[] = preferredProviders; + if (resolved.variant === "ctc") { + const { session, providersUsed } = await this.createSessionWithFallback( + resolved.ctcModelPath!, + preferredProviders, + ); + this.ctcSession = session; + this.ctcOutputName = + session.outputNames.find((name) => /logprob|logits/i.test(name)) || + session.outputNames[0] || + null; + + logger.transcription.info("Initialized local Parakeet model", { + modelId: requestedId, + variant: resolved.variant, + modelPath: resolved.ctcModelPath, + vocabPath: resolved.vocabPath, + executionProviders: providersUsed, + outputName: this.ctcOutputName, + featureSize: this.featureSize, + }); + } else { + const encoderResult = await this.createSessionWithFallback( + resolved.encoderModelPath!, + preferredProviders, + ); + const decoderResult = await this.createSessionWithFallback( + resolved.decoderJointModelPath!, + preferredProviders, + ); + + this.tdtEncoderSession = encoderResult.session; + this.tdtDecoderJointSession = decoderResult.session; + this.ctcOutputName = null; + + logger.transcription.info("Initialized local Parakeet model", { + modelId: requestedId, + variant: resolved.variant, + encoderModelPath: resolved.encoderModelPath, + decoderJointModelPath: resolved.decoderJointModelPath, + vocabPath: resolved.vocabPath, + executionProviders: { + encoder: encoderResult.providersUsed, + decoder: decoderResult.providersUsed, + }, + featureSize: this.featureSize, + maxTokensPerStep: this.maxTokensPerStep, + }); + } + + this.currentModelId = requestedId; + this.currentModelVariant = resolved.variant; + } + private async createSessionWithFallback( + modelPath: string, + preferredProviders: readonly string[], + ): Promise<{ session: ort.InferenceSession; providersUsed: readonly string[] }> { try { - session = await ort.InferenceSession.create(modelPath, { + const session = await ort.InferenceSession.create(modelPath, { executionProviders: [...preferredProviders], }); + return { session, providersUsed: preferredProviders }; } catch (error) { if (preferredProviders.length > 1) { logger.transcription.warn( "Parakeet preferred execution provider unavailable, falling back to CPU", { requestedProviders: preferredProviders, + modelPath, error: error instanceof Error ? error.message : String(error), }, ); - session = await ort.InferenceSession.create(modelPath, { + const session = await ort.InferenceSession.create(modelPath, { executionProviders: ["cpu"], }); - providersUsed = ["cpu"]; - } else { - throw error; + return { session, providersUsed: ["cpu"] }; } + throw error; } + } - this.session = session; - this.outputName = - session.outputNames.find((name) => /logprob|logits/i.test(name)) || - session.outputNames[0] || - null; - this.vocabulary = await loadParakeetVocabulary(vocabPath); - this.currentModelId = requestedId; - - logger.transcription.info("Initialized local Parakeet model", { - modelId: requestedId, - modelPath, - vocabPath, - executionProviders: providersUsed, - outputName: this.outputName, - }); + private async releaseSessions(): Promise { + if (this.ctcSession) { + await this.ctcSession.release(); + this.ctcSession = null; + } + if (this.tdtEncoderSession) { + await this.tdtEncoderSession.release(); + this.tdtEncoderSession = null; + } + if (this.tdtDecoderJointSession) { + await this.tdtDecoderJointSession.release(); + this.tdtDecoderJointSession = null; + } } private async resolveSelectedParakeetModelId( @@ -241,10 +542,7 @@ export class ParakeetProvider implements TranscriptionProvider { ): Promise { const selectedId = requestedModelId || (await this.modelService.getSelectedModel()); if (!selectedId) { - throw new AppError( - "No speech model selected", - ErrorCodes.MODEL_MISSING, - ); + throw new AppError("No speech model selected", ErrorCodes.MODEL_MISSING); } if (!selectedId.startsWith("parakeet-")) { @@ -257,10 +555,7 @@ export class ParakeetProvider implements TranscriptionProvider { return selectedId; } - private async resolveModelPaths(modelId: string): Promise<{ - modelPath: string; - vocabPath: string; - }> { + private async resolveModelPaths(modelId: string): Promise { const downloadedModels = await this.modelService.getDownloadedModels(); const downloaded = downloadedModels[modelId]; @@ -271,23 +566,244 @@ export class ParakeetProvider implements TranscriptionProvider { ); } - const modelPath = downloaded.localPath; - const modelDir = path.dirname(modelPath); - + const modelDir = path.dirname(downloaded.localPath); const localFiles = downloaded.originalModel && typeof downloaded.originalModel === "object" && Array.isArray((downloaded.originalModel as { localFiles?: unknown }).localFiles) - ? ((downloaded.originalModel as { localFiles: unknown[] }).localFiles.filter( + ? (downloaded.originalModel as { localFiles: unknown[] }).localFiles.filter( (value): value is string => typeof value === "string", - ) as string[]) - : []; + ) + : [downloaded.localPath]; + + const findFile = (pattern: RegExp): string | undefined => + localFiles.find((filePath) => pattern.test(path.basename(filePath))); const vocabPath = - localFiles.find((filePath) => filePath.endsWith("vocab.txt")) || - path.join(modelDir, "vocab.txt"); + findFile(/^vocab\.txt$/i) || path.join(modelDir, "vocab.txt"); + + const decoderJointModelPath = findFile(/^decoder_joint-model(?:\.int8)?\.onnx$/i); + const encoderModelPath = findFile(/^encoder-model(?:\.int8)?\.onnx$/i); + + if (decoderJointModelPath || encoderModelPath) { + if (!decoderJointModelPath || !encoderModelPath) { + throw new AppError( + `Parakeet TDT artifacts are incomplete for ${modelId}`, + ErrorCodes.MODEL_MISSING, + ); + } + + const configPath = + findFile(/^config\.json$/i) || path.join(modelDir, "config.json"); + + return { + variant: "tdt", + encoderModelPath, + decoderJointModelPath, + vocabPath, + configPath, + }; + } + + const ctcModelPath = + findFile(/^model(?:\.int8)?\.onnx$/i) || downloaded.localPath; + + return { + variant: "ctc", + ctcModelPath, + vocabPath, + }; + } + + private async loadModelConfig(configPath?: string): Promise { + if (!configPath) { + return {}; + } + + try { + const content = await fs.readFile(configPath, "utf8"); + const parsed = JSON.parse(content) as ParakeetModelConfig; + return parsed && typeof parsed === "object" ? parsed : {}; + } catch { + return {}; + } + } + + private createEncoderAccessor(encoderTensor: ort.Tensor): EncoderAccessor { + const dims = encoderTensor.dims; + const data = this.toFloat32Array(encoderTensor.data); + + if (dims.length !== 3 || dims[0] !== 1) { + throw new AppError( + `Unexpected Parakeet encoder output dims: ${dims.join("x")}`, + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const dim1 = Number(dims[1]); + const dim2 = Number(dims[2]); + + if (!Number.isFinite(dim1) || !Number.isFinite(dim2) || dim1 <= 0 || dim2 <= 0) { + throw new AppError( + `Invalid Parakeet encoder output dims: ${dims.join("x")}`, + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + // Typical NeMo output shape is [1, hidden, time]. + const hiddenFirst = dim1 >= dim2; + const hiddenSize = hiddenFirst ? dim1 : dim2; + const timeSteps = hiddenFirst ? dim2 : dim1; + + return { + hiddenSize, + timeSteps, + at: (timeStep: number): Float32Array => { + const step = Math.max(0, Math.min(timeSteps - 1, timeStep)); + const vector = new Float32Array(hiddenSize); + + if (hiddenFirst) { + for (let h = 0; h < hiddenSize; h++) { + vector[h] = data[h * timeSteps + step] ?? 0; + } + } else { + for (let h = 0; h < hiddenSize; h++) { + vector[h] = data[step * hiddenSize + h] ?? 0; + } + } + + return vector; + }, + }; + } + + private createInitialTdtState( + session: ort.InferenceSession, + state1InputName: string, + state2InputName: string, + fallbackHiddenSize: number, + ): TdtDecoderState { + const state1Shape = this.getInputTensorShape(session, state1InputName); + const state2Shape = this.getInputTensorShape(session, state2InputName); + + const layers = this.dimToNumber(state1Shape?.[0], 2); + const hidden1 = this.dimToNumber(state1Shape?.[2], fallbackHiddenSize); + const hidden2 = this.dimToNumber(state2Shape?.[2], hidden1); + + return { + state1: new ort.Tensor("float32", new Float32Array(layers * hidden1), [layers, 1, hidden1]), + state2: new ort.Tensor("float32", new Float32Array(layers * hidden2), [layers, 1, hidden2]), + }; + } + + private getInputTensorShape( + session: ort.InferenceSession, + inputName: string, + ): ReadonlyArray | null { + const index = session.inputNames.indexOf(inputName); + if (index < 0) { + return null; + } + + const metadata = session.inputMetadata[index]; + if (!metadata || !metadata.isTensor) { + return null; + } + + return metadata.shape; + } + + private dimToNumber(value: number | string | undefined, fallback: number): number { + if (typeof value === "number" && Number.isFinite(value) && value > 0) { + return value; + } + return fallback; + } + + private createIntegerTensorForInput( + session: ort.InferenceSession, + inputName: string, + values: number[], + dims: readonly number[], + ): ort.Tensor { + const inputType = this.getInputTensorType(session, inputName); + if (inputType === "int32") { + return new ort.Tensor( + "int32", + Int32Array.from(values.map((value) => Math.trunc(value))), + Array.from(dims), + ); + } + + // Default to int64 for NeMo/Parakeet paths that use long tensors. + return new ort.Tensor( + "int64", + BigInt64Array.from( + values.map((value) => BigInt(Math.trunc(value))), + ), + Array.from(dims), + ); + } + + private getInputTensorType( + session: ort.InferenceSession, + inputName: string, + ): ort.Tensor.Type | null { + const index = session.inputNames.indexOf(inputName); + if (index < 0) { + return null; + } + + const metadata = session.inputMetadata[index]; + if (!metadata || !metadata.isTensor) { + return null; + } + + return metadata.type; + } + + private findName( + names: readonly string[], + patterns: RegExp[], + fallbackIndex = 0, + ): string { + for (const pattern of patterns) { + const matched = names.find((name) => pattern.test(name)); + if (matched) { + return matched; + } + } + return names[fallbackIndex] || names[0] || ""; + } + + private toFloat32Array(data: ort.Tensor["data"]): Float32Array { + if (data instanceof Float32Array) { + return data; + } + return Float32Array.from(data as ArrayLike); + } + + private toBigInt64Array(data: ort.Tensor["data"]): BigInt64Array { + if (data instanceof BigInt64Array) { + return data; + } + const values = Array.from(data as ArrayLike, (value) => BigInt(Math.trunc(value))); + return BigInt64Array.from(values); + } + + private argmax(values: Float32Array, start: number, length: number): number { + let bestIndex = 0; + let bestValue = -Number.MAX_VALUE; + + for (let i = 0; i < length; i++) { + const value = values[start + i] ?? -Number.MAX_VALUE; + if (value > bestValue) { + bestValue = value; + bestIndex = i; + } + } - return { modelPath, vocabPath }; + return bestIndex; } private shouldTranscribe(): boolean { diff --git a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts index 8d8115cf..e92132a0 100644 --- a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts +++ b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts @@ -6,7 +6,7 @@ const WIN_LENGTH = 400; const HOP_LENGTH = 160; const PREEMPHASIS = 0.97; const LOG_ZERO_GUARD = Math.pow(2, -24); -const N_MELS = 80; +const DEFAULT_N_MELS = 80; const F_MIN = 0; const F_MAX = SAMPLE_RATE / 2; const DECODE_SPACE_PATTERN = /^\s|\s\B|(\s)\b/g; @@ -39,27 +39,27 @@ function buildCenteredHannWindow(): Float32Array { return window; } -function buildMelFilterBank(): Float32Array[] { +function buildMelFilterBank(numMels: number): Float32Array[] { const numBins = Math.floor(N_FFT / 2) + 1; const fbanks: Float32Array[] = Array.from( - { length: N_MELS }, + { length: numMels }, () => new Float32Array(numBins), ); const minMel = hzToMel(F_MIN); const maxMel = hzToMel(F_MAX); - const melPoints = new Float64Array(N_MELS + 2); + const melPoints = new Float64Array(numMels + 2); for (let i = 0; i < melPoints.length; i++) { - melPoints[i] = minMel + ((maxMel - minMel) * i) / (N_MELS + 1); + melPoints[i] = minMel + ((maxMel - minMel) * i) / (numMels + 1); } - const bins = new Int32Array(N_MELS + 2); + const bins = new Int32Array(numMels + 2); for (let i = 0; i < bins.length; i++) { const hz = melToHz(melPoints[i]); bins[i] = Math.floor(((N_FFT + 1) * hz) / SAMPLE_RATE); } - for (let m = 1; m <= N_MELS; m++) { + for (let m = 1; m <= numMels; m++) { const left = bins[m - 1]; const center = bins[m]; const right = bins[m + 1]; @@ -165,11 +165,17 @@ function fftInPlace( } export class ParakeetFeatureExtractor { + private readonly nMels: number; private readonly window = buildCenteredHannWindow(); - private readonly melBanks = buildMelFilterBank(); + private readonly melBanks: Float32Array[]; private readonly bitReverse = createBitReverseTable(N_FFT); private readonly twiddle = createTwiddleTables(N_FFT); + constructor(nMels = DEFAULT_N_MELS) { + this.nMels = nMels; + this.melBanks = buildMelFilterBank(nMels); + } + extract(audioData: Float32Array): ParakeetFeatures { const preemphasized = new Float32Array(audioData.length); if (audioData.length > 0) { @@ -185,7 +191,7 @@ export class ParakeetFeatureExtractor { const frameCount = Math.max(1, Math.floor((padded.length - N_FFT) / HOP_LENGTH) + 1); const featuresLength = Math.max(1, Math.floor(audioData.length / HOP_LENGTH)); - const logMel = new Float32Array(frameCount * N_MELS); + const logMel = new Float32Array(frameCount * this.nMels); const real = new Float32Array(N_FFT); const imag = new Float32Array(N_FFT); @@ -206,30 +212,30 @@ export class ParakeetFeatureExtractor { this.twiddle.sin, ); - for (let m = 0; m < N_MELS; m++) { + for (let m = 0; m < this.nMels; m++) { const bank = this.melBanks[m]; let energy = 0; for (let k = 0; k < bank.length; k++) { const power = real[k] * real[k] + imag[k] * imag[k]; energy += power * bank[k]; } - logMel[frame * N_MELS + m] = Math.log(energy + LOG_ZERO_GUARD); + logMel[frame * this.nMels + m] = Math.log(energy + LOG_ZERO_GUARD); } } const validFrames = Math.min(featuresLength, frameCount); - const normalized = new Float32Array(N_MELS * frameCount); + const normalized = new Float32Array(this.nMels * frameCount); - for (let m = 0; m < N_MELS; m++) { + for (let m = 0; m < this.nMels; m++) { let mean = 0; for (let f = 0; f < validFrames; f++) { - mean += logMel[f * N_MELS + m]; + mean += logMel[f * this.nMels + m]; } mean /= validFrames; let variance = 0; for (let f = 0; f < validFrames; f++) { - const delta = logMel[f * N_MELS + m] - mean; + const delta = logMel[f * this.nMels + m] - mean; variance += delta * delta; } const denom = Math.max(validFrames - 1, 1); @@ -238,18 +244,28 @@ export class ParakeetFeatureExtractor { const invStd = 1 / (Math.sqrt(variance) + 1e-5); for (let f = 0; f < frameCount; f++) { normalized[m * frameCount + f] = - f < validFrames ? (logMel[f * N_MELS + m] - mean) * invStd : 0; + f < validFrames ? (logMel[f * this.nMels + m] - mean) * invStd : 0; } } return { inputFeatures: normalized, - inputShape: [1, N_MELS, frameCount], + inputShape: [1, this.nMels, frameCount], featuresLength: validFrames, }; } } +export function decodeParakeetTokens(tokenIds: number[], vocab: string[]): string { + const text = tokenIds + .map((id) => vocab[id] ?? "") + .filter((token) => token && !token.startsWith("<|") && token !== "") + .join(""); + return text.replace(DECODE_SPACE_PATTERN, (_match, capturedSpace) => { + return capturedSpace ? " " : ""; + }); +} + export async function loadParakeetVocabulary( vocabPath: string, ): Promise { @@ -326,8 +342,5 @@ export function decodeParakeetCtc( prevId = bestId; } - const text = tokenIds.map((id) => vocab[id] ?? "").join(""); - return text.replace(DECODE_SPACE_PATTERN, (_match, capturedSpace) => { - return capturedSpace ? " " : ""; - }); + return decodeParakeetTokens(tokenIds, vocab); } diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index 23812905..42309272 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -62,6 +62,7 @@ class ModelService extends EventEmitter { private settingsService: SettingsService; private readonly localSpeechPreference = [ "whisper-large-v3-turbo", + "parakeet-tdt-0.6b-v3-int8", "parakeet-ctc-0.6b-int8", "whisper-large-v3", "whisper-medium", diff --git a/apps/desktop/src/trpc/routers/models.ts b/apps/desktop/src/trpc/routers/models.ts index e65f5669..23a38b12 100644 --- a/apps/desktop/src/trpc/routers/models.ts +++ b/apps/desktop/src/trpc/routers/models.ts @@ -43,6 +43,12 @@ export const modelsRouter = createRouter({ // Include setup field from available model metadata return { ...downloaded, + // Always prefer current manifest metadata for display fields. + name: m.name, + size: m.sizeFormatted, + description: m.description, + speed: m.speed, + accuracy: m.accuracy, setup: m.setup, } as Model & { setup: "offline" | "cloud" }; } From 1d7907b4c1deecf84b8bee7f1c8d48ce6933a296 Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Sat, 14 Feb 2026 20:37:43 +0000 Subject: [PATCH 3/8] feat(desktop): use nemo onnx preprocessor for parakeet tdt --- .../transcription/parakeet-provider.ts | 118 +++++++++++++++++- 1 file changed, 117 insertions(+), 1 deletion(-) diff --git a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts index b9b2e7a2..3f0ff36f 100644 --- a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -26,6 +26,7 @@ interface ResolvedParakeetPaths { ctcModelPath?: string; encoderModelPath?: string; decoderJointModelPath?: string; + nemoPreprocessorPath?: string; vocabPath: string; configPath?: string; } @@ -50,6 +51,7 @@ export class ParakeetProvider implements TranscriptionProvider { readonly name = "parakeet-local"; private ctcSession: ort.InferenceSession | null = null; + private tdtPreprocessorSession: ort.InferenceSession | null = null; private tdtEncoderSession: ort.InferenceSession | null = null; private tdtDecoderJointSession: ort.InferenceSession | null = null; @@ -155,7 +157,7 @@ export class ParakeetProvider implements TranscriptionProvider { variant: this.currentModelVariant, }); - const features = this.featureExtractor.extract(speechAudio); + const features = await this.extractFeatures(speechAudio); if (this.currentModelVariant === "tdt") { return this.transcribeTdt(features); @@ -174,6 +176,89 @@ export class ParakeetProvider implements TranscriptionProvider { } } + private async extractFeatures(audioData: Float32Array): Promise { + if (this.currentModelVariant !== "tdt" || !this.tdtPreprocessorSession) { + return this.featureExtractor.extract(audioData); + } + + try { + const waveformInputName = this.findName( + this.tdtPreprocessorSession.inputNames, + [/waveforms/i, /audio/i], + 0, + ); + const waveformLengthInputName = this.findName( + this.tdtPreprocessorSession.inputNames, + [/waveforms_lens/i, /length/i], + 1, + ); + const featuresOutputName = this.findName( + this.tdtPreprocessorSession.outputNames, + [/^features$/i, /mel/i], + 0, + ); + const featuresLengthOutputName = this.findName( + this.tdtPreprocessorSession.outputNames, + [/features_lens/i, /length/i], + 1, + ); + + const waveformTensor = new ort.Tensor("float32", audioData, [1, audioData.length]); + const waveformLengthTensor = this.createIntegerTensorForInput( + this.tdtPreprocessorSession, + waveformLengthInputName, + [audioData.length], + [1], + ); + + const preprocessorResults = await this.tdtPreprocessorSession.run({ + [waveformInputName]: waveformTensor, + [waveformLengthInputName]: waveformLengthTensor, + }); + + const featuresTensor = preprocessorResults[featuresOutputName] as ort.Tensor; + const featuresLengthTensor = preprocessorResults[featuresLengthOutputName] as ort.Tensor; + const featuresData = this.toFloat32Array(featuresTensor.data); + const dims = featuresTensor.dims; + + if (dims.length !== 3 || dims[0] !== 1) { + throw new AppError( + `Unexpected Parakeet preprocessor output dims: ${dims.join("x")}`, + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const featuresSize = Number(dims[1]); + const frameCount = Number(dims[2]); + if (!Number.isFinite(featuresSize) || !Number.isFinite(frameCount) || featuresSize <= 0 || frameCount <= 0) { + throw new AppError( + `Invalid Parakeet preprocessor output dims: ${dims.join("x")}`, + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const featureLengths = this.toBigInt64Array(featuresLengthTensor.data); + const featuresLength = Math.max( + 1, + Math.min(frameCount, Number(featureLengths[0] ?? BigInt(frameCount))), + ); + + return { + inputFeatures: featuresData, + inputShape: [1, featuresSize, frameCount], + featuresLength, + }; + } catch (error) { + logger.transcription.warn( + "Parakeet TDT ONNX preprocessor failed, falling back to JS feature extraction", + { + error: error instanceof Error ? error.message : String(error), + }, + ); + return this.featureExtractor.extract(audioData); + } + } + private async transcribeCtc(features: ParakeetFeatures): Promise { if (!this.ctcSession || !this.vocabulary || !this.ctcOutputName) { throw new AppError( @@ -462,6 +547,16 @@ export class ParakeetProvider implements TranscriptionProvider { featureSize: this.featureSize, }); } else { + let preprocessorProviders: readonly string[] | null = null; + if (resolved.nemoPreprocessorPath) { + const preprocessorResult = await this.createSessionWithFallback( + resolved.nemoPreprocessorPath, + preferredProviders, + ); + this.tdtPreprocessorSession = preprocessorResult.session; + preprocessorProviders = preprocessorResult.providersUsed; + } + const encoderResult = await this.createSessionWithFallback( resolved.encoderModelPath!, preferredProviders, @@ -478,10 +573,12 @@ export class ParakeetProvider implements TranscriptionProvider { logger.transcription.info("Initialized local Parakeet model", { modelId: requestedId, variant: resolved.variant, + nemoPreprocessorPath: resolved.nemoPreprocessorPath || null, encoderModelPath: resolved.encoderModelPath, decoderJointModelPath: resolved.decoderJointModelPath, vocabPath: resolved.vocabPath, executionProviders: { + preprocessor: preprocessorProviders, encoder: encoderResult.providersUsed, decoder: decoderResult.providersUsed, }, @@ -527,6 +624,10 @@ export class ParakeetProvider implements TranscriptionProvider { await this.ctcSession.release(); this.ctcSession = null; } + if (this.tdtPreprocessorSession) { + await this.tdtPreprocessorSession.release(); + this.tdtPreprocessorSession = null; + } if (this.tdtEncoderSession) { await this.tdtEncoderSession.release(); this.tdtEncoderSession = null; @@ -595,11 +696,17 @@ export class ParakeetProvider implements TranscriptionProvider { const configPath = findFile(/^config\.json$/i) || path.join(modelDir, "config.json"); + const nemoPathCandidate = + findFile(/^nemo\d+\.onnx$/i) || path.join(modelDir, "nemo128.onnx"); + const nemoPreprocessorPath = await this.fileExists(nemoPathCandidate) + ? nemoPathCandidate + : undefined; return { variant: "tdt", encoderModelPath, decoderJointModelPath, + nemoPreprocessorPath, vocabPath, configPath, }; @@ -615,6 +722,15 @@ export class ParakeetProvider implements TranscriptionProvider { }; } + private async fileExists(filePath: string): Promise { + try { + await fs.access(filePath); + return true; + } catch { + return false; + } + } + private async loadModelConfig(configPath?: string): Promise { if (!configPath) { return {}; From 5527641e931f3bb8a1ec991ea280f14b31290111 Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Fri, 20 Mar 2026 20:21:29 +0000 Subject: [PATCH 4/8] feat(desktop): remove legacy parakeet ctc support --- apps/desktop/src/constants/models.ts | 51 --- .../transcription/parakeet-provider.ts | 342 ++++++++---------- .../utils/parakeet-feature-extractor.ts | 63 +--- apps/desktop/src/services/model-service.ts | 79 ++-- 4 files changed, 224 insertions(+), 311 deletions(-) diff --git a/apps/desktop/src/constants/models.ts b/apps/desktop/src/constants/models.ts index cf1bc351..9d7ab031 100644 --- a/apps/desktop/src/constants/models.ts +++ b/apps/desktop/src/constants/models.ts @@ -157,57 +157,6 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ provider: "Amical Cloud", providerIcon: "/assets/logo.svg", }, - { - id: "parakeet-ctc-0.6b-int8", - name: "NVIDIA Parakeet CTC 0.6B", - type: "whisper", - description: - "CTC speech model optimized for fast English transcription with ONNX Runtime.", - checksum: - "3cfe22e14a7adf70b7b4ab33109a6fad4d7ca61821ca1f37168dc9b3d04b963b", - filename: "model.int8.onnx", - artifacts: [ - { - filename: "model.int8.onnx", - downloadUrl: - "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/model.int8.onnx", - checksum: - "3cfe22e14a7adf70b7b4ab33109a6fad4d7ca61821ca1f37168dc9b3d04b963b", - size: 653436437, - }, - { - filename: "vocab.txt", - downloadUrl: - "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/vocab.txt", - }, - ], - downloadUrl: - "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx/resolve/main/model.int8.onnx", - size: 653436437, - sizeFormatted: "~623 MB", - modelSize: "~623 MB", - features: [ - { - icon: "bolt", - tooltip: "Fast local transcription", - }, - { - icon: "gauge", - tooltip: "DirectML/CPU ONNX runtime", - }, - { - icon: "languages", - tooltip: "English-first CTC vocabulary", - }, - ], - speed: 4.6, - accuracy: 4.2, - setup: "offline", - runtime: "parakeet-onnx", - provider: "NVIDIA", - providerIcon: "/icons/models/nvidia.svg", - sourceUrl: "https://huggingface.co/istupakov/parakeet-ctc-0.6b-onnx", - }, { id: "parakeet-tdt-0.6b-v3-int8", name: "NVIDIA Parakeet TDT 0.6B v3", diff --git a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts index 3f0ff36f..1c84937e 100644 --- a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -12,20 +12,15 @@ import { AppError, ErrorCodes } from "../../../types/error"; import { extractSpeechFromVad } from "../../utils/vad-audio-filter"; import { ParakeetFeatureExtractor, - decodeParakeetCtc, decodeParakeetTokens, loadParakeetVocabulary, ParakeetVocabulary, ParakeetFeatures, } from "../../utils/parakeet-feature-extractor"; -type ParakeetModelVariant = "ctc" | "tdt"; - interface ResolvedParakeetPaths { - variant: ParakeetModelVariant; - ctcModelPath?: string; - encoderModelPath?: string; - decoderJointModelPath?: string; + encoderModelPath: string; + decoderJointModelPath: string; nemoPreprocessorPath?: string; vocabPath: string; configPath?: string; @@ -50,14 +45,11 @@ interface TdtDecoderState { export class ParakeetProvider implements TranscriptionProvider { readonly name = "parakeet-local"; - private ctcSession: ort.InferenceSession | null = null; private tdtPreprocessorSession: ort.InferenceSession | null = null; private tdtEncoderSession: ort.InferenceSession | null = null; private tdtDecoderJointSession: ort.InferenceSession | null = null; - private ctcOutputName: string | null = null; private currentModelId: string | null = null; - private currentModelVariant: ParakeetModelVariant | null = null; private vocabulary: ParakeetVocabulary | null = null; private frameBuffer: Float32Array[] = []; @@ -118,16 +110,14 @@ export class ParakeetProvider implements TranscriptionProvider { async dispose(): Promise { await this.releaseSessions(); - this.ctcOutputName = null; this.currentModelId = null; - this.currentModelVariant = null; this.vocabulary = null; this.reset(); } private async doTranscription(_context: TranscribeContext): Promise { try { - if (!this.vocabulary || !this.currentModelVariant) { + if (!this.vocabulary) { throw new AppError( "Parakeet model is not initialized", ErrorCodes.WORKER_INITIALIZATION_FAILED, @@ -154,16 +144,10 @@ export class ParakeetProvider implements TranscriptionProvider { before: rawAudio.length, after: speechAudio.length, segments: segments.length, - variant: this.currentModelVariant, }); const features = await this.extractFeatures(speechAudio); - - if (this.currentModelVariant === "tdt") { - return this.transcribeTdt(features); - } - - return this.transcribeCtc(features); + return this.transcribeTdt(features); } catch (error) { logger.transcription.error("Parakeet transcription failed", { error }); if (error instanceof AppError) { @@ -176,8 +160,10 @@ export class ParakeetProvider implements TranscriptionProvider { } } - private async extractFeatures(audioData: Float32Array): Promise { - if (this.currentModelVariant !== "tdt" || !this.tdtPreprocessorSession) { + private async extractFeatures( + audioData: Float32Array, + ): Promise { + if (!this.tdtPreprocessorSession) { return this.featureExtractor.extract(audioData); } @@ -203,7 +189,10 @@ export class ParakeetProvider implements TranscriptionProvider { 1, ); - const waveformTensor = new ort.Tensor("float32", audioData, [1, audioData.length]); + const waveformTensor = new ort.Tensor("float32", audioData, [ + 1, + audioData.length, + ]); const waveformLengthTensor = this.createIntegerTensorForInput( this.tdtPreprocessorSession, waveformLengthInputName, @@ -216,8 +205,12 @@ export class ParakeetProvider implements TranscriptionProvider { [waveformLengthInputName]: waveformLengthTensor, }); - const featuresTensor = preprocessorResults[featuresOutputName] as ort.Tensor; - const featuresLengthTensor = preprocessorResults[featuresLengthOutputName] as ort.Tensor; + const featuresTensor = preprocessorResults[ + featuresOutputName + ] as ort.Tensor; + const featuresLengthTensor = preprocessorResults[ + featuresLengthOutputName + ] as ort.Tensor; const featuresData = this.toFloat32Array(featuresTensor.data); const dims = featuresTensor.dims; @@ -230,7 +223,12 @@ export class ParakeetProvider implements TranscriptionProvider { const featuresSize = Number(dims[1]); const frameCount = Number(dims[2]); - if (!Number.isFinite(featuresSize) || !Number.isFinite(frameCount) || featuresSize <= 0 || frameCount <= 0) { + if ( + !Number.isFinite(featuresSize) || + !Number.isFinite(frameCount) || + featuresSize <= 0 || + frameCount <= 0 + ) { throw new AppError( `Invalid Parakeet preprocessor output dims: ${dims.join("x")}`, ErrorCodes.WORKER_INITIALIZATION_FAILED, @@ -259,48 +257,12 @@ export class ParakeetProvider implements TranscriptionProvider { } } - private async transcribeCtc(features: ParakeetFeatures): Promise { - if (!this.ctcSession || !this.vocabulary || !this.ctcOutputName) { - throw new AppError( - "Parakeet CTC model is not initialized", - ErrorCodes.WORKER_INITIALIZATION_FAILED, - ); - } - - const inputName = this.findName(this.ctcSession.inputNames, [/audio_signal/i], 0); - const lengthName = this.findName(this.ctcSession.inputNames, [/length/i], 1); - - const inputTensor = new ort.Tensor( - "float32", - features.inputFeatures, - features.inputShape, - ); - const lengthTensor = this.createIntegerTensorForInput( - this.ctcSession, - lengthName, - [features.featuresLength], - [1], - ); - - const results = await this.ctcSession.run({ - [inputName]: inputTensor, - [lengthName]: lengthTensor, - }); - - const logitsTensor = results[this.ctcOutputName] as ort.Tensor; - const logits = this.toFloat32Array(logitsTensor.data); - - return decodeParakeetCtc( - logits, - logitsTensor.dims, - this.vocabulary.tokens, - this.vocabulary.blankTokenId, - features.featuresLength, - ); - } - private async transcribeTdt(features: ParakeetFeatures): Promise { - if (!this.tdtEncoderSession || !this.tdtDecoderJointSession || !this.vocabulary) { + if ( + !this.tdtEncoderSession || + !this.tdtDecoderJointSession || + !this.vocabulary + ) { throw new AppError( "Parakeet TDT model is not initialized", ErrorCodes.WORKER_INITIALIZATION_FAILED, @@ -353,7 +315,10 @@ export class ParakeetProvider implements TranscriptionProvider { const encodedLengths = this.toBigInt64Array(encodedLengthTensor.data); const encodedLength = Math.max( 1, - Math.min(accessor.timeSteps, Number(encodedLengths[0] ?? BigInt(accessor.timeSteps))), + Math.min( + accessor.timeSteps, + Number(encodedLengths[0] ?? BigInt(accessor.timeSteps)), + ), ); const decoderOutputName = this.findName( @@ -414,13 +379,14 @@ export class ParakeetProvider implements TranscriptionProvider { const guardLimit = encodedLength * Math.max(8, this.maxTokensPerStep * 4); while (t < encodedLength && guard++ < guardLimit) { - const encoderStepTensor = new ort.Tensor( - "float32", - accessor.at(t), - [1, accessor.hiddenSize, 1], - ); + const encoderStepTensor = new ort.Tensor("float32", accessor.at(t), [ + 1, + accessor.hiddenSize, + 1, + ]); - const lastTokenId = tokenIds.length > 0 ? tokenIds[tokenIds.length - 1] : blankTokenId; + const lastTokenId = + tokenIds.length > 0 ? tokenIds[tokenIds.length - 1] : blankTokenId; const targetsTensor = this.createIntegerTensorForInput( this.tdtDecoderJointSession, decoderTargetsInputName, @@ -472,7 +438,10 @@ export class ParakeetProvider implements TranscriptionProvider { if (stepCount > 0) { t += stepCount; emittedTokens = 0; - } else if (token === blankTokenId || emittedTokens >= this.maxTokensPerStep) { + } else if ( + token === blankTokenId || + emittedTokens >= this.maxTokensPerStep + ) { t += 1; emittedTokens = 0; } @@ -493,10 +462,8 @@ export class ParakeetProvider implements TranscriptionProvider { if ( this.vocabulary && this.currentModelId === requestedId && - ((this.currentModelVariant === "ctc" && this.ctcSession) || - (this.currentModelVariant === "tdt" && - this.tdtEncoderSession && - this.tdtDecoderJointSession)) + this.tdtEncoderSession && + this.tdtDecoderJointSession ) { return; } @@ -506,11 +473,7 @@ export class ParakeetProvider implements TranscriptionProvider { const config = await this.loadModelConfig(resolved.configPath); this.featureSize = - typeof config.features_size === "number" - ? config.features_size - : resolved.variant === "tdt" - ? 128 - : 80; + typeof config.features_size === "number" ? config.features_size : 128; this.maxTokensPerStep = typeof config.max_tokens_per_step === "number" ? config.max_tokens_per_step @@ -526,75 +489,53 @@ export class ParakeetProvider implements TranscriptionProvider { ? (["coreml", "cpu"] as const) : (["cpu"] as const); - if (resolved.variant === "ctc") { - const { session, providersUsed } = await this.createSessionWithFallback( - resolved.ctcModelPath!, + let preprocessorProviders: readonly string[] | null = null; + if (resolved.nemoPreprocessorPath) { + const preprocessorResult = await this.createSessionWithFallback( + resolved.nemoPreprocessorPath, preferredProviders, ); - this.ctcSession = session; - this.ctcOutputName = - session.outputNames.find((name) => /logprob|logits/i.test(name)) || - session.outputNames[0] || - null; - - logger.transcription.info("Initialized local Parakeet model", { - modelId: requestedId, - variant: resolved.variant, - modelPath: resolved.ctcModelPath, - vocabPath: resolved.vocabPath, - executionProviders: providersUsed, - outputName: this.ctcOutputName, - featureSize: this.featureSize, - }); - } else { - let preprocessorProviders: readonly string[] | null = null; - if (resolved.nemoPreprocessorPath) { - const preprocessorResult = await this.createSessionWithFallback( - resolved.nemoPreprocessorPath, - preferredProviders, - ); - this.tdtPreprocessorSession = preprocessorResult.session; - preprocessorProviders = preprocessorResult.providersUsed; - } + this.tdtPreprocessorSession = preprocessorResult.session; + preprocessorProviders = preprocessorResult.providersUsed; + } - const encoderResult = await this.createSessionWithFallback( - resolved.encoderModelPath!, - preferredProviders, - ); - const decoderResult = await this.createSessionWithFallback( - resolved.decoderJointModelPath!, - preferredProviders, - ); + const encoderResult = await this.createSessionWithFallback( + resolved.encoderModelPath, + preferredProviders, + ); + const decoderResult = await this.createSessionWithFallback( + resolved.decoderJointModelPath, + preferredProviders, + ); - this.tdtEncoderSession = encoderResult.session; - this.tdtDecoderJointSession = decoderResult.session; - this.ctcOutputName = null; - - logger.transcription.info("Initialized local Parakeet model", { - modelId: requestedId, - variant: resolved.variant, - nemoPreprocessorPath: resolved.nemoPreprocessorPath || null, - encoderModelPath: resolved.encoderModelPath, - decoderJointModelPath: resolved.decoderJointModelPath, - vocabPath: resolved.vocabPath, - executionProviders: { - preprocessor: preprocessorProviders, - encoder: encoderResult.providersUsed, - decoder: decoderResult.providersUsed, - }, - featureSize: this.featureSize, - maxTokensPerStep: this.maxTokensPerStep, - }); - } + this.tdtEncoderSession = encoderResult.session; + this.tdtDecoderJointSession = decoderResult.session; + + logger.transcription.info("Initialized local Parakeet model", { + modelId: requestedId, + nemoPreprocessorPath: resolved.nemoPreprocessorPath || null, + encoderModelPath: resolved.encoderModelPath, + decoderJointModelPath: resolved.decoderJointModelPath, + vocabPath: resolved.vocabPath, + executionProviders: { + preprocessor: preprocessorProviders, + encoder: encoderResult.providersUsed, + decoder: decoderResult.providersUsed, + }, + featureSize: this.featureSize, + maxTokensPerStep: this.maxTokensPerStep, + }); this.currentModelId = requestedId; - this.currentModelVariant = resolved.variant; } private async createSessionWithFallback( modelPath: string, preferredProviders: readonly string[], - ): Promise<{ session: ort.InferenceSession; providersUsed: readonly string[] }> { + ): Promise<{ + session: ort.InferenceSession; + providersUsed: readonly string[]; + }> { try { const session = await ort.InferenceSession.create(modelPath, { executionProviders: [...preferredProviders], @@ -620,10 +561,6 @@ export class ParakeetProvider implements TranscriptionProvider { } private async releaseSessions(): Promise { - if (this.ctcSession) { - await this.ctcSession.release(); - this.ctcSession = null; - } if (this.tdtPreprocessorSession) { await this.tdtPreprocessorSession.release(); this.tdtPreprocessorSession = null; @@ -641,7 +578,8 @@ export class ParakeetProvider implements TranscriptionProvider { private async resolveSelectedParakeetModelId( requestedModelId?: string, ): Promise { - const selectedId = requestedModelId || (await this.modelService.getSelectedModel()); + const selectedId = + requestedModelId || (await this.modelService.getSelectedModel()); if (!selectedId) { throw new AppError("No speech model selected", ErrorCodes.MODEL_MISSING); } @@ -656,7 +594,9 @@ export class ParakeetProvider implements TranscriptionProvider { return selectedId; } - private async resolveModelPaths(modelId: string): Promise { + private async resolveModelPaths( + modelId: string, + ): Promise { const downloadedModels = await this.modelService.getDownloadedModels(); const downloaded = downloadedModels[modelId]; @@ -671,8 +611,12 @@ export class ParakeetProvider implements TranscriptionProvider { const localFiles = downloaded.originalModel && typeof downloaded.originalModel === "object" && - Array.isArray((downloaded.originalModel as { localFiles?: unknown }).localFiles) - ? (downloaded.originalModel as { localFiles: unknown[] }).localFiles.filter( + Array.isArray( + (downloaded.originalModel as { localFiles?: unknown }).localFiles, + ) + ? ( + downloaded.originalModel as { localFiles: unknown[] } + ).localFiles.filter( (value): value is string => typeof value === "string", ) : [downloaded.localPath]; @@ -683,42 +627,42 @@ export class ParakeetProvider implements TranscriptionProvider { const vocabPath = findFile(/^vocab\.txt$/i) || path.join(modelDir, "vocab.txt"); - const decoderJointModelPath = findFile(/^decoder_joint-model(?:\.int8)?\.onnx$/i); - const encoderModelPath = findFile(/^encoder-model(?:\.int8)?\.onnx$/i); - - if (decoderJointModelPath || encoderModelPath) { - if (!decoderJointModelPath || !encoderModelPath) { - throw new AppError( - `Parakeet TDT artifacts are incomplete for ${modelId}`, - ErrorCodes.MODEL_MISSING, - ); - } - - const configPath = - findFile(/^config\.json$/i) || path.join(modelDir, "config.json"); - const nemoPathCandidate = - findFile(/^nemo\d+\.onnx$/i) || path.join(modelDir, "nemo128.onnx"); - const nemoPreprocessorPath = await this.fileExists(nemoPathCandidate) - ? nemoPathCandidate - : undefined; - - return { - variant: "tdt", - encoderModelPath, - decoderJointModelPath, - nemoPreprocessorPath, - vocabPath, - configPath, - }; + const decoderJointPathCandidate = + findFile(/^decoder_joint-model(?:\.int8)?\.onnx$/i) || + path.join(modelDir, "decoder_joint-model.int8.onnx"); + const encoderPathCandidate = + findFile(/^encoder-model(?:\.int8)?\.onnx$/i) || + path.join(modelDir, "encoder-model.int8.onnx"); + const decoderJointModelPath = (await this.fileExists( + decoderJointPathCandidate, + )) + ? decoderJointPathCandidate + : undefined; + const encoderModelPath = (await this.fileExists(encoderPathCandidate)) + ? encoderPathCandidate + : undefined; + + if (!decoderJointModelPath || !encoderModelPath) { + throw new AppError( + `Parakeet TDT artifacts are incomplete for ${modelId}`, + ErrorCodes.MODEL_MISSING, + ); } - const ctcModelPath = - findFile(/^model(?:\.int8)?\.onnx$/i) || downloaded.localPath; + const configPath = + findFile(/^config\.json$/i) || path.join(modelDir, "config.json"); + const nemoPathCandidate = + findFile(/^nemo\d+\.onnx$/i) || path.join(modelDir, "nemo128.onnx"); + const nemoPreprocessorPath = (await this.fileExists(nemoPathCandidate)) + ? nemoPathCandidate + : undefined; return { - variant: "ctc", - ctcModelPath, + encoderModelPath, + decoderJointModelPath, + nemoPreprocessorPath, vocabPath, + configPath, }; } @@ -731,7 +675,9 @@ export class ParakeetProvider implements TranscriptionProvider { } } - private async loadModelConfig(configPath?: string): Promise { + private async loadModelConfig( + configPath?: string, + ): Promise { if (!configPath) { return {}; } @@ -759,7 +705,12 @@ export class ParakeetProvider implements TranscriptionProvider { const dim1 = Number(dims[1]); const dim2 = Number(dims[2]); - if (!Number.isFinite(dim1) || !Number.isFinite(dim2) || dim1 <= 0 || dim2 <= 0) { + if ( + !Number.isFinite(dim1) || + !Number.isFinite(dim2) || + dim1 <= 0 || + dim2 <= 0 + ) { throw new AppError( `Invalid Parakeet encoder output dims: ${dims.join("x")}`, ErrorCodes.WORKER_INITIALIZATION_FAILED, @@ -807,8 +758,16 @@ export class ParakeetProvider implements TranscriptionProvider { const hidden2 = this.dimToNumber(state2Shape?.[2], hidden1); return { - state1: new ort.Tensor("float32", new Float32Array(layers * hidden1), [layers, 1, hidden1]), - state2: new ort.Tensor("float32", new Float32Array(layers * hidden2), [layers, 1, hidden2]), + state1: new ort.Tensor("float32", new Float32Array(layers * hidden1), [ + layers, + 1, + hidden1, + ]), + state2: new ort.Tensor("float32", new Float32Array(layers * hidden2), [ + layers, + 1, + hidden2, + ]), }; } @@ -829,7 +788,10 @@ export class ParakeetProvider implements TranscriptionProvider { return metadata.shape; } - private dimToNumber(value: number | string | undefined, fallback: number): number { + private dimToNumber( + value: number | string | undefined, + fallback: number, + ): number { if (typeof value === "number" && Number.isFinite(value) && value > 0) { return value; } @@ -854,9 +816,7 @@ export class ParakeetProvider implements TranscriptionProvider { // Default to int64 for NeMo/Parakeet paths that use long tensors. return new ort.Tensor( "int64", - BigInt64Array.from( - values.map((value) => BigInt(Math.trunc(value))), - ), + BigInt64Array.from(values.map((value) => BigInt(Math.trunc(value)))), Array.from(dims), ); } @@ -903,7 +863,9 @@ export class ParakeetProvider implements TranscriptionProvider { if (data instanceof BigInt64Array) { return data; } - const values = Array.from(data as ArrayLike, (value) => BigInt(Math.trunc(value))); + const values = Array.from(data as ArrayLike, (value) => + BigInt(Math.trunc(value)), + ); return BigInt64Array.from(values); } diff --git a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts index e92132a0..1dce4fbc 100644 --- a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts +++ b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts @@ -34,7 +34,8 @@ function buildCenteredHannWindow(): Float32Array { const window = new Float32Array(N_FFT); const pad = (N_FFT - WIN_LENGTH) / 2; for (let i = 0; i < WIN_LENGTH; i++) { - window[pad + i] = 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (WIN_LENGTH - 1)); + window[pad + i] = + 0.5 - 0.5 * Math.cos((2 * Math.PI * i) / (WIN_LENGTH - 1)); } return window; } @@ -188,8 +189,14 @@ export class ParakeetFeatureExtractor { const padded = new Float32Array(preemphasized.length + N_FFT); padded.set(preemphasized, N_FFT / 2); - const frameCount = Math.max(1, Math.floor((padded.length - N_FFT) / HOP_LENGTH) + 1); - const featuresLength = Math.max(1, Math.floor(audioData.length / HOP_LENGTH)); + const frameCount = Math.max( + 1, + Math.floor((padded.length - N_FFT) / HOP_LENGTH) + 1, + ); + const featuresLength = Math.max( + 1, + Math.floor(audioData.length / HOP_LENGTH), + ); const logMel = new Float32Array(frameCount * this.nMels); const real = new Float32Array(N_FFT); @@ -256,7 +263,10 @@ export class ParakeetFeatureExtractor { } } -export function decodeParakeetTokens(tokenIds: number[], vocab: string[]): string { +export function decodeParakeetTokens( + tokenIds: number[], + vocab: string[], +): string { const text = tokenIds .map((id) => vocab[id] ?? "") .filter((token) => token && !token.startsWith("<|") && token !== "") @@ -299,48 +309,3 @@ export async function loadParakeetVocabulary( blankTokenId: blankTokenId >= 0 ? blankTokenId : tokens.length - 1, }; } - -export function decodeParakeetCtc( - logits: Float32Array, - dims: readonly number[], - vocab: string[], - blankTokenId: number, - requestedLength: number, -): string { - if (dims.length !== 3 || dims[0] !== 1) { - return ""; - } - - const [_, dim1, dim2] = dims; - const vocabSize = vocab.length; - - const timeFirst = dim2 === vocabSize; - const timeSteps = timeFirst ? dim1 : dim2; - const availableLength = Math.min(requestedLength, timeSteps); - - const tokenIds: number[] = []; - let prevId = blankTokenId; - - for (let t = 0; t < availableLength; t++) { - let bestId = 0; - let bestLogit = -Number.MAX_VALUE; - - for (let v = 0; v < vocabSize; v++) { - const index = timeFirst - ? t * vocabSize + v - : v * timeSteps + t; - const value = logits[index] ?? -Number.MAX_VALUE; - if (value > bestLogit) { - bestLogit = value; - bestId = v; - } - } - - if (bestId !== blankTokenId && bestId !== prevId) { - tokenIds.push(bestId); - } - prevId = bestId; - } - - return decodeParakeetTokens(tokenIds, vocab); -} diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index 42309272..cd751b14 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -56,14 +56,16 @@ interface ModelManagerEvents { ) => void; } +const LEGACY_PARAKEET_CTC_MODEL_ID = "parakeet-ctc-0.6b-int8"; +const PARAKEET_TDT_MODEL_ID = "parakeet-tdt-0.6b-v3-int8"; + class ModelService extends EventEmitter { private state: ModelManagerState; private modelsDirectory: string; private settingsService: SettingsService; private readonly localSpeechPreference = [ "whisper-large-v3-turbo", - "parakeet-tdt-0.6b-v3-int8", - "parakeet-ctc-0.6b-int8", + PARAKEET_TDT_MODEL_ID, "whisper-large-v3", "whisper-medium", "whisper-small", @@ -116,18 +118,18 @@ class ModelService extends EventEmitter { async initialize(): Promise { try { // Sync local speech models with filesystem - const whisperModelsData = AVAILABLE_MODELS - .filter((model) => model.setup === "offline" && !!model.filename) - .map((model) => ({ - id: model.id, - name: model.name, - description: model.description, - size: model.sizeFormatted, - checksum: model.checksum, - speed: model.speed, - accuracy: model.accuracy, - filename: model.filename, - })); + const whisperModelsData = AVAILABLE_MODELS.filter( + (model) => model.setup === "offline" && !!model.filename, + ).map((model) => ({ + id: model.id, + name: model.name, + description: model.description, + size: model.sizeFormatted, + checksum: model.checksum, + speed: model.speed, + accuracy: model.accuracy, + filename: model.filename, + })); const syncResult = await syncLocalWhisperModels( this.modelsDirectory, @@ -141,7 +143,7 @@ class ModelService extends EventEmitter { }); // Restore selected model from settings and validate availability - const savedSelection = await this.settingsService.getDefaultSpeechModel(); + const savedSelection = await this.getSelectedModel(); if (savedSelection) { // Validate the saved selection is still available @@ -546,8 +548,7 @@ class ModelService extends EventEmitter { // Auto-select if this is the first model const allDownloadedModels = await this.getValidDownloadedModels(); const downloadedModelCount = Object.keys(allDownloadedModels).length; - const currentSelection = - await this.settingsService.getDefaultSpeechModel(); + const currentSelection = await this.getSelectedModel(); if (downloadedModelCount === 1 && !currentSelection) { await this.applySpeechModelSelection( @@ -611,7 +612,7 @@ class ModelService extends EventEmitter { } // Check if this is the selected model BEFORE deletion - const currentSelection = await this.settingsService.getDefaultSpeechModel(); + const currentSelection = await this.getSelectedModel(); const wasSelected = currentSelection === modelId; // Delete file @@ -725,7 +726,40 @@ class ModelService extends EventEmitter { // Get currently selected model for transcription async getSelectedModel(): Promise { - return (await this.settingsService.getDefaultSpeechModel()) || null; + const selectedModelId = + (await this.settingsService.getDefaultSpeechModel()) || null; + return this.normalizeLegacySpeechModelSelection(selectedModelId); + } + + private async normalizeLegacySpeechModelSelection( + modelId: string | null, + ): Promise { + if (modelId !== LEGACY_PARAKEET_CTC_MODEL_ID) { + return modelId; + } + + const downloadedModels = await this.getValidDownloadedModels(); + const fallbackModelIds = Object.keys(downloadedModels).filter( + (downloadedModelId) => downloadedModelId !== LEGACY_PARAKEET_CTC_MODEL_ID, + ); + const fallbackModelId = downloadedModels[PARAKEET_TDT_MODEL_ID] + ? PARAKEET_TDT_MODEL_ID + : fallbackModelIds.length > 0 + ? this.pickPreferredLocalModelId(fallbackModelIds) + : null; + + await this.applySpeechModelSelection( + fallbackModelId, + fallbackModelId ? "auto-after-deletion" : "cleared", + modelId, + ); + + logger.main.info("Migrated legacy Parakeet CTC speech selection", { + from: modelId, + to: fallbackModelId, + }); + + return fallbackModelId; } private async syncFormatterConfigForSpeechChange( @@ -806,6 +840,10 @@ class ModelService extends EventEmitter { // Set selected model for transcription async setSelectedModel(modelId: string | null): Promise { + if (modelId === LEGACY_PARAKEET_CTC_MODEL_ID) { + modelId = PARAKEET_TDT_MODEL_ID; + } + const oldModelId = await this.getSelectedModel(); // If setting to a specific model, validate it exists @@ -1185,8 +1223,7 @@ class ModelService extends EventEmitter { */ async validateAndClearInvalidDefaults(): Promise { // Check default speech model - const defaultSpeechModel = - await this.settingsService.getDefaultSpeechModel(); + const defaultSpeechModel = await this.getSelectedModel(); if (defaultSpeechModel) { const availableModel = AVAILABLE_MODELS.find( (m) => m.id === defaultSpeechModel, From 8fbb0d76c283f6208f85c246cc7209b4e6d90163 Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Fri, 20 Mar 2026 20:22:49 +0000 Subject: [PATCH 5/8] chore(desktop): restore default native helper script --- apps/desktop/package.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/desktop/package.json b/apps/desktop/package.json index 7717504d..0adea6f2 100644 --- a/apps/desktop/package.json +++ b/apps/desktop/package.json @@ -44,7 +44,7 @@ "build:types": "pnpm --filter @amical/types build", "build:swift-helper": "pnpm --filter @amical/swift-helper build", "build:windows-helper": "pnpm --filter @amical/windows-helper build", - "build:native-helper": "node -e \"const { execSync } = require('node:child_process'); const cmd = process.platform === 'darwin' ? 'pnpm run build:swift-helper' : process.platform === 'win32' ? 'pnpm run build:windows-helper' : 'echo No native helpers'; execSync(cmd, { stdio: 'inherit' });\"", + "build:native-helper": "node -p \"process.platform === 'darwin' ? 'build:swift-helper' : process.platform === 'win32' ? 'build:windows-helper' : 'echo No native helpers'\" | xargs pnpm run", "dev": "pnpm start", "download-node": "tsx scripts/download-node-binaries.ts", "download-node:all": "tsx scripts/download-node-binaries.ts --all" From addd346e87c82e90ebceb2ffa8b81f35f7e2a864 Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Fri, 20 Mar 2026 21:06:48 +0000 Subject: [PATCH 6/8] fix(desktop): require complete local model bundles --- apps/desktop/src/db/models.ts | 92 ++++++++++-- apps/desktop/src/services/model-service.ts | 65 +++++++-- apps/desktop/src/trpc/routers/models.ts | 2 +- .../tests/services/model-validity.test.ts | 133 ++++++++++++++++++ 4 files changed, 269 insertions(+), 23 deletions(-) create mode 100644 apps/desktop/tests/services/model-validity.test.ts diff --git a/apps/desktop/src/db/models.ts b/apps/desktop/src/db/models.ts index d8c1f627..8731d287 100644 --- a/apps/desktop/src/db/models.ts +++ b/apps/desktop/src/db/models.ts @@ -172,6 +172,9 @@ export async function syncLocalWhisperModels( speed: number; accuracy: number; filename: string; + artifacts?: Array<{ + filename: string; + }>; }>, ): Promise<{ added: number; updated: number; removed: number }> { const fs = await import("fs"); @@ -188,33 +191,96 @@ export async function syncLocalWhisperModels( // Map available models by ID for easy lookup // (we already have them indexed by ID, so we don't need this map) + const resolveRequiredLocalFiles = ( + model: (typeof availableModels)[number], + ) => { + const requiredFilenames = + model.artifacts && model.artifacts.length > 0 + ? model.artifacts.map((artifact) => artifact.filename) + : [model.filename]; + + const resolvedFiles = requiredFilenames + .map((filename) => { + const candidatePaths = [ + path.join(modelsDirectory, filename), + path.join(modelsDirectory, model.id, filename), + ]; + + return candidatePaths.find((candidatePath) => + fs.existsSync(candidatePath), + ); + }) + .filter((filePath): filePath is string => !!filePath); + + return resolvedFiles.length === requiredFilenames.length + ? resolvedFiles + : null; + }; + // Process each available model for (const model of availableModels) { - const candidatePaths = [ - path.join(modelsDirectory, model.filename), - path.join(modelsDirectory, model.id, model.filename), - ]; + const resolvedFiles = resolveRequiredLocalFiles(model); const filePath = - candidatePaths.find((candidatePath) => fs.existsSync(candidatePath)) || - candidatePaths[1]; - const fileExists = fs.existsSync(filePath); + resolvedFiles?.find( + (resolvedFilePath) => + path.basename(resolvedFilePath) === model.filename, + ) || path.join(modelsDirectory, model.id, model.filename); + const fileExists = !!resolvedFiles; const existingRecord = existingModelMap.get(model.id); if (fileExists) { - // File exists on disk - const stats = fs.statSync(filePath); + const sizeBytes = resolvedFiles.reduce( + (sum, resolvedFilePath) => sum + fs.statSync(resolvedFilePath).size, + 0, + ); + const existingLocalFiles = + existingRecord?.originalModel && + typeof existingRecord.originalModel === "object" && + !Array.isArray(existingRecord.originalModel) && + Array.isArray( + ( + existingRecord.originalModel as { + localFiles?: unknown; + } + ).localFiles, + ) + ? ( + existingRecord.originalModel as { + localFiles: unknown[]; + } + ).localFiles.filter( + (value): value is string => typeof value === "string", + ) + : []; + const localFilesChanged = + existingLocalFiles.length !== resolvedFiles.length || + existingLocalFiles.some( + (existingLocalFile, index) => + existingLocalFile !== resolvedFiles[index], + ); + const originalModel = + existingRecord?.originalModel && + typeof existingRecord.originalModel === "object" && + !Array.isArray(existingRecord.originalModel) + ? { + ...existingRecord.originalModel, + localFiles: resolvedFiles, + } + : { localFiles: resolvedFiles }; if (existingRecord) { // Update existing record if needed if ( existingRecord.localPath !== filePath || - existingRecord.sizeBytes !== stats.size + existingRecord.sizeBytes !== sizeBytes || + localFilesChanged ) { await upsertModel({ ...existingRecord, localPath: filePath, - sizeBytes: stats.size, + sizeBytes, downloadedAt: existingRecord.downloadedAt || new Date(), + originalModel, }); updated++; } @@ -231,10 +297,10 @@ export async function syncLocalWhisperModels( speed: model.speed, accuracy: model.accuracy, localPath: filePath, - sizeBytes: stats.size, + sizeBytes, downloadedAt: new Date(), context: null, - originalModel: null, + originalModel, }); added++; } diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index cd751b14..fa73fe00 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -14,7 +14,6 @@ import { getModelsByProvider, getDownloadedWhisperModels, removeModel, - modelExists, syncLocalWhisperModels, getAllModels, syncModelsForProvider, @@ -129,6 +128,9 @@ class ModelService extends EventEmitter { speed: model.speed, accuracy: model.accuracy, filename: model.filename, + artifacts: model.artifacts?.map((artifact) => ({ + filename: artifact.filename, + })), })); const syncResult = await syncLocalWhisperModels( @@ -319,17 +321,24 @@ class ModelService extends EventEmitter { return record; } - // Get only valid downloaded models (files that exist on disk) - // Since we sync on init and only store downloaded models, all models in DB are valid + // Get only valid downloaded models (all required artifacts exist on disk) async getValidDownloadedModels(): Promise> { - return this.getDownloadedModels(); + const downloadedModels = await this.getDownloadedModels(); + const validModels: Record = {}; + + for (const model of Object.values(downloadedModels)) { + if (this.isDownloadedModelValid(model)) { + validModels[model.id] = model; + } + } + + return validModels; } // Check if a model is downloaded - // Since we only store downloaded models, just check if it exists in DB async isModelDownloaded(modelId: string): Promise { - const models = await getModelsByProvider("local-whisper"); - return models.some((m) => m.id === modelId); + const downloadedModels = await this.getValidDownloadedModels(); + return !!downloadedModels[modelId]; } // Get download progress for a model @@ -1229,10 +1238,11 @@ class ModelService extends EventEmitter { (m) => m.id === defaultSpeechModel, ); const isAmicalModel = availableModel?.setup === "cloud"; - const existsInDb = await modelExists("local-whisper", defaultSpeechModel); + const validDownloadedModels = await this.getValidDownloadedModels(); + const existsLocally = !!validDownloadedModels[defaultSpeechModel]; // Amical cloud models are always valid; local models must exist in DB - if (!isAmicalModel && !existsInDb) { + if (!isAmicalModel && !existsLocally) { logger.main.info("Clearing invalid default speech model", { modelId: defaultSpeechModel, }); @@ -1294,6 +1304,43 @@ class ModelService extends EventEmitter { } } } + + private isDownloadedModelValid(model: DBModel): boolean { + if (!model.localPath) { + return false; + } + + const availableModel = AVAILABLE_MODELS.find( + (available) => available.id === model.id && available.setup === "offline", + ); + if (!availableModel) { + return false; + } + + const modelDirectory = path.dirname(model.localPath); + const recordedLocalFiles = + model.originalModel && + typeof model.originalModel === "object" && + Array.isArray( + (model.originalModel as { localFiles?: unknown }).localFiles, + ) + ? (model.originalModel as { localFiles: unknown[] }).localFiles.filter( + (value): value is string => typeof value === "string", + ) + : []; + const requiredFilenames = + availableModel.artifacts && availableModel.artifacts.length > 0 + ? availableModel.artifacts.map((artifact) => artifact.filename) + : [availableModel.filename]; + + return requiredFilenames.every((filename) => { + const recordedMatch = recordedLocalFiles.find( + (localFile) => path.basename(localFile) === filename, + ); + const resolvedPath = recordedMatch || path.join(modelDirectory, filename); + return fs.existsSync(resolvedPath); + }); + } } export { ModelService }; diff --git a/apps/desktop/src/trpc/routers/models.ts b/apps/desktop/src/trpc/routers/models.ts index 23a38b12..1bce8a25 100644 --- a/apps/desktop/src/trpc/routers/models.ts +++ b/apps/desktop/src/trpc/routers/models.ts @@ -30,7 +30,7 @@ export const modelsRouter = createRouter({ // Return all available whisper models as Model type // We need to convert from AvailableWhisperModel to Model format const availableModels = modelService.getAvailableModels(); - const downloadedModels = await modelService.getDownloadedModels(); + const downloadedModels = await modelService.getValidDownloadedModels(); // Check authentication status for cloud model filtering const authService = ctx.serviceManager.getService("authService"); diff --git a/apps/desktop/tests/services/model-validity.test.ts b/apps/desktop/tests/services/model-validity.test.ts new file mode 100644 index 00000000..a3e12e97 --- /dev/null +++ b/apps/desktop/tests/services/model-validity.test.ts @@ -0,0 +1,133 @@ +import path from "node:path"; +import fs from "fs-extra"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { + syncLocalWhisperModels, + getDownloadedWhisperModels, + upsertModel, +} from "../../src/db/models"; +import { ModelService } from "../../src/services/model-service"; +import { createTestDatabase, type TestDatabase } from "../helpers/test-db"; +import { TEST_USER_DATA_PATH } from "../helpers/electron-mocks"; +import { setTestDatabase } from "../setup"; + +const PARAKEET_MODEL = { + id: "parakeet-tdt-0.6b-v3-int8", + name: "NVIDIA Parakeet TDT 0.6B v3", + description: "Test model", + size: "~640 MB", + speed: 4.3, + accuracy: 4.6, + filename: "encoder-model.int8.onnx", + artifacts: [ + { filename: "encoder-model.int8.onnx" }, + { filename: "decoder_joint-model.int8.onnx" }, + { filename: "nemo128.onnx" }, + { filename: "vocab.txt" }, + { filename: "config.json" }, + ], +} as const; + +describe("Model validity", () => { + let testDb: TestDatabase; + let modelsDirectory: string; + + beforeEach(async () => { + testDb = await createTestDatabase({ name: "model-validity-test" }); + setTestDatabase(testDb.db); + modelsDirectory = path.join(TEST_USER_DATA_PATH, "models"); + await fs.emptyDir(modelsDirectory); + }); + + afterEach(async () => { + await fs.emptyDir(modelsDirectory); + await testDb.close(); + }); + + it("does not sync a Parakeet install unless all artifacts are present", async () => { + const modelDirectory = path.join(modelsDirectory, PARAKEET_MODEL.id); + await fs.ensureDir(modelDirectory); + await fs.writeFile( + path.join(modelDirectory, "encoder-model.int8.onnx"), + "encoder", + ); + await fs.writeFile( + path.join(modelDirectory, "decoder_joint-model.int8.onnx"), + "decoder", + ); + + const result = await syncLocalWhisperModels(modelsDirectory, [ + PARAKEET_MODEL, + ]); + const downloadedModels = await getDownloadedWhisperModels(); + + expect(result.added).toBe(0); + expect(downloadedModels).toHaveLength(0); + }); + + it("syncs a complete Parakeet install with all local files recorded", async () => { + const modelDirectory = path.join(modelsDirectory, PARAKEET_MODEL.id); + await fs.ensureDir(modelDirectory); + + for (const artifact of PARAKEET_MODEL.artifacts) { + await fs.writeFile( + path.join(modelDirectory, artifact.filename), + artifact.filename, + ); + } + + const result = await syncLocalWhisperModels(modelsDirectory, [ + PARAKEET_MODEL, + ]); + const downloadedModels = await getDownloadedWhisperModels(); + + expect(result.added).toBe(1); + expect(downloadedModels).toHaveLength(1); + expect(downloadedModels[0].localPath).toBe( + path.join(modelDirectory, PARAKEET_MODEL.filename), + ); + expect(downloadedModels[0].sizeBytes).toBe( + PARAKEET_MODEL.artifacts.reduce( + (sum, artifact) => sum + Buffer.byteLength(artifact.filename), + 0, + ), + ); + expect(downloadedModels[0].originalModel).toEqual({ + localFiles: PARAKEET_MODEL.artifacts.map((artifact) => + path.join(modelDirectory, artifact.filename), + ), + }); + }); + + it("treats partial Parakeet bundles as not downloaded at runtime", async () => { + const modelDirectory = path.join(modelsDirectory, PARAKEET_MODEL.id); + await fs.ensureDir(modelDirectory); + const localPath = path.join(modelDirectory, PARAKEET_MODEL.filename); + await fs.writeFile(localPath, "encoder"); + await fs.writeFile(path.join(modelDirectory, "vocab.txt"), "vocab"); + + await upsertModel({ + id: PARAKEET_MODEL.id, + provider: "local-whisper", + name: PARAKEET_MODEL.name, + type: "speech", + size: PARAKEET_MODEL.size, + description: PARAKEET_MODEL.description, + localPath, + sizeBytes: 11, + checksum: null, + downloadedAt: new Date(), + originalModel: { + localFiles: [localPath, path.join(modelDirectory, "vocab.txt")], + }, + speed: PARAKEET_MODEL.speed, + accuracy: PARAKEET_MODEL.accuracy, + context: null, + }); + + const modelService = new ModelService({} as never); + + expect(await modelService.isModelDownloaded(PARAKEET_MODEL.id)).toBe(false); + expect(await modelService.getValidDownloadedModels()).toEqual({}); + }); +}); From 1652647baca49874f1b7cf1d8d77d11c24bea765 Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Sat, 21 Mar 2026 06:12:48 +0000 Subject: [PATCH 7/8] fix(desktop): harden parakeet transcription init --- .../transcription/parakeet-provider.ts | 110 ++++++++++++------ .../utils/parakeet-feature-extractor.ts | 13 ++- apps/desktop/src/services/model-service.ts | 108 ++++++++++++----- 3 files changed, 165 insertions(+), 66 deletions(-) diff --git a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts index 1c84937e..ba89ac59 100644 --- a/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -124,9 +124,9 @@ export class ParakeetProvider implements TranscriptionProvider { ); } + const bufferedFrames = [...this.frameBuffer]; const vadProbs = [...this.frameBufferSpeechProbabilities]; - const rawAudio = this.aggregateFrames(); - this.reset(); + const rawAudio = this.aggregateFrames(bufferedFrames); const { audio: speechAudio, segments } = extractSpeechFromVad( rawAudio, @@ -137,6 +137,7 @@ export class ParakeetProvider implements TranscriptionProvider { logger.transcription.debug( "Skipping Parakeet transcription - no speech detected by VAD filter", ); + this.reset(); return ""; } @@ -147,7 +148,9 @@ export class ParakeetProvider implements TranscriptionProvider { }); const features = await this.extractFeatures(speechAudio); - return this.transcribeTdt(features); + const transcript = await this.transcribeTdt(features); + this.reset(); + return transcript; } catch (error) { logger.transcription.error("Parakeet transcription failed", { error }); if (error instanceof AppError) { @@ -469,18 +472,16 @@ export class ParakeetProvider implements TranscriptionProvider { } const resolved = await this.resolveModelPaths(requestedId); - await this.releaseSessions(); - const config = await this.loadModelConfig(resolved.configPath); - this.featureSize = + const newFeatureSize = typeof config.features_size === "number" ? config.features_size : 128; - this.maxTokensPerStep = + const newMaxTokensPerStep = typeof config.max_tokens_per_step === "number" ? config.max_tokens_per_step : 10; - this.featureExtractor = new ParakeetFeatureExtractor(this.featureSize); + const newFeatureExtractor = new ParakeetFeatureExtractor(newFeatureSize); - this.vocabulary = await loadParakeetVocabulary(resolved.vocabPath); + const newVocabulary = await loadParakeetVocabulary(resolved.vocabPath); const preferredProviders = process.platform === "win32" @@ -490,26 +491,58 @@ export class ParakeetProvider implements TranscriptionProvider { : (["cpu"] as const); let preprocessorProviders: readonly string[] | null = null; - if (resolved.nemoPreprocessorPath) { - const preprocessorResult = await this.createSessionWithFallback( - resolved.nemoPreprocessorPath, + let newPreprocessorSession: ort.InferenceSession | null = null; + let newEncoderSession: ort.InferenceSession | null = null; + let newDecoderSession: ort.InferenceSession | null = null; + let encoderProviders: readonly string[] = ["cpu"]; + let decoderProviders: readonly string[] = ["cpu"]; + + try { + if (resolved.nemoPreprocessorPath) { + const preprocessorResult = await this.createSessionWithFallback( + resolved.nemoPreprocessorPath, + preferredProviders, + ); + newPreprocessorSession = preprocessorResult.session; + preprocessorProviders = preprocessorResult.providersUsed; + } + + const encoderResult = await this.createSessionWithFallback( + resolved.encoderModelPath, preferredProviders, ); - this.tdtPreprocessorSession = preprocessorResult.session; - preprocessorProviders = preprocessorResult.providersUsed; - } - - const encoderResult = await this.createSessionWithFallback( - resolved.encoderModelPath, - preferredProviders, - ); - const decoderResult = await this.createSessionWithFallback( - resolved.decoderJointModelPath, - preferredProviders, - ); + newEncoderSession = encoderResult.session; + encoderProviders = encoderResult.providersUsed; - this.tdtEncoderSession = encoderResult.session; - this.tdtDecoderJointSession = decoderResult.session; + const decoderResult = await this.createSessionWithFallback( + resolved.decoderJointModelPath, + preferredProviders, + ); + newDecoderSession = decoderResult.session; + decoderProviders = decoderResult.providersUsed; + + await this.releaseSessions(); + + this.featureSize = newFeatureSize; + this.maxTokensPerStep = newMaxTokensPerStep; + this.featureExtractor = newFeatureExtractor; + this.vocabulary = newVocabulary; + this.tdtPreprocessorSession = newPreprocessorSession; + this.tdtEncoderSession = newEncoderSession; + this.tdtDecoderJointSession = newDecoderSession; + this.currentModelId = requestedId; + + newPreprocessorSession = null; + newEncoderSession = null; + newDecoderSession = null; + } catch (error) { + await Promise.allSettled([ + this.releaseSession(newPreprocessorSession), + this.releaseSession(newEncoderSession), + this.releaseSession(newDecoderSession), + ]); + throw error; + } logger.transcription.info("Initialized local Parakeet model", { modelId: requestedId, @@ -519,14 +552,12 @@ export class ParakeetProvider implements TranscriptionProvider { vocabPath: resolved.vocabPath, executionProviders: { preprocessor: preprocessorProviders, - encoder: encoderResult.providersUsed, - decoder: decoderResult.providersUsed, + encoder: encoderProviders, + decoder: decoderProviders, }, featureSize: this.featureSize, maxTokensPerStep: this.maxTokensPerStep, }); - - this.currentModelId = requestedId; } private async createSessionWithFallback( @@ -575,6 +606,14 @@ export class ParakeetProvider implements TranscriptionProvider { } } + private async releaseSession( + session: ort.InferenceSession | null, + ): Promise { + if (session) { + await session.release(); + } + } + private async resolveSelectedParakeetModelId( requestedModelId?: string, ): Promise { @@ -905,15 +944,14 @@ export class ParakeetProvider implements TranscriptionProvider { return false; } - private aggregateFrames(): Float32Array { - const totalLength = this.frameBuffer.reduce( - (sum, frame) => sum + frame.length, - 0, - ); + private aggregateFrames( + frames: readonly Float32Array[] = this.frameBuffer, + ): Float32Array { + const totalLength = frames.reduce((sum, frame) => sum + frame.length, 0); const aggregated = new Float32Array(totalLength); let offset = 0; - for (const frame of this.frameBuffer) { + for (const frame of frames) { aggregated.set(frame, offset); offset += frame.length; } diff --git a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts index 1dce4fbc..a59e072b 100644 --- a/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts +++ b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts @@ -296,16 +296,27 @@ export async function loadParakeetVocabulary( tokensById.set(id, token); } + if (tokensById.size === 0) { + throw new Error(`Parakeet vocabulary is empty or invalid: ${vocabPath}`); + } + const maxId = Math.max(...tokensById.keys()); const tokens: string[] = Array.from({ length: maxId + 1 }, () => ""); for (const [id, token] of tokensById.entries()) { tokens[id] = token; } + if (tokens.length === 0) { + throw new Error(`Parakeet vocabulary produced no tokens: ${vocabPath}`); + } + const blankTokenId = tokens.findIndex((token) => token === ""); + if (blankTokenId < 0) { + throw new Error(`Parakeet vocabulary is missing : ${vocabPath}`); + } return { tokens, - blankTokenId: blankTokenId >= 0 ? blankTokenId : tokens.length - 1, + blankTokenId, }; } diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index fa73fe00..bb8b53b0 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -444,47 +444,97 @@ class ModelService extends EventEmitter { } const fileStream = fs.createWriteStream(artifactPath); + let fileStreamError: Error | null = null; + const onFileStreamError = (error: Error) => { + fileStreamError = error; + }; + fileStream.on("error", onFileStreamError); const reader = response.body?.getReader(); if (!reader) { + fileStream.destroy(); throw new Error(`Failed to read ${artifact.filename}`); } - while (true) { - const { done, value } = await reader.read(); - if (done) break; + let writeCompleted = false; + try { + while (true) { + if (fileStreamError) { + throw fileStreamError; + } + + const { done, value } = await reader.read(); + if (done) break; + + if (abortController.signal.aborted) { + fileStream.close(); + if (fs.existsSync(artifactPath)) { + fs.unlinkSync(artifactPath); + } + throw new Error("Download cancelled"); + } + + const canContinue = fileStream.write(value); + if (!canContinue) { + await new Promise((resolve, reject) => { + const onDrain = () => { + fileStream.off("error", onDrainError); + resolve(); + }; + const onDrainError = (error: Error) => { + fileStream.off("drain", onDrain); + reject(error); + }; + + fileStream.once("drain", onDrain); + fileStream.once("error", onDrainError); + }); + } + + if (fileStreamError) { + throw fileStreamError; + } - if (abortController.signal.aborted) { - fileStream.close(); - if (fs.existsSync(artifactPath)) { - fs.unlinkSync(artifactPath); + bytesDownloaded += value.length; + progress.bytesDownloaded = bytesDownloaded; + progress.progress = + progress.totalBytes > 0 + ? Math.round((bytesDownloaded / progress.totalBytes) * 100) + : 0; + + const progressPercent = progress.progress; + if ( + progressPercent - lastProgressEmit >= 1 || + bytesDownloaded - + (lastProgressEmit * progress.totalBytes) / 100 >= + 1024 * 1024 + ) { + this.emit("download-progress", modelId, { ...progress }); + lastProgressEmit = progressPercent; } - throw new Error("Download cancelled"); } - fileStream.write(value); - bytesDownloaded += value.length; - progress.bytesDownloaded = bytesDownloaded; - progress.progress = - progress.totalBytes > 0 - ? Math.round((bytesDownloaded / progress.totalBytes) * 100) - : 0; - - const progressPercent = progress.progress; - if ( - progressPercent - lastProgressEmit >= 1 || - bytesDownloaded - (lastProgressEmit * progress.totalBytes) / 100 >= - 1024 * 1024 - ) { - this.emit("download-progress", modelId, { ...progress }); - lastProgressEmit = progressPercent; + await new Promise((resolve, reject) => { + const onFinish = () => { + fileStream.off("error", onFinishError); + resolve(); + }; + const onFinishError = (error: Error) => { + fileStream.off("finish", onFinish); + reject(error); + }; + + fileStream.once("finish", onFinish); + fileStream.once("error", onFinishError); + fileStream.end(); + }); + writeCompleted = true; + } finally { + fileStream.off("error", onFileStreamError); + if (!writeCompleted && !fileStream.closed) { + fileStream.destroy(); } } - await new Promise((resolve, reject) => { - fileStream.end(() => resolve()); - fileStream.on("error", reject); - }); - if (artifact.checksum) { const fileChecksum = await this.calculateFileChecksum( artifactPath, From 70345d5aadcdc8fa49408361e4d351f486f8c7bf Mon Sep 17 00:00:00 2001 From: Oliver Atkinson Date: Sat, 21 Mar 2026 06:59:57 +0000 Subject: [PATCH 8/8] fix(desktop): wait for stream close before cleanup --- apps/desktop/src/services/model-service.ts | 64 +++++++++++++++------- 1 file changed, 43 insertions(+), 21 deletions(-) diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index bb8b53b0..5eb43fc4 100644 --- a/apps/desktop/src/services/model-service.ts +++ b/apps/desktop/src/services/model-service.ts @@ -449,13 +449,47 @@ class ModelService extends EventEmitter { fileStreamError = error; }; fileStream.on("error", onFileStreamError); + const awaitFileStreamClose = async ( + closeAction: () => void, + preserveOriginalError = false, + ): Promise => { + if (fileStream.closed) { + if (!preserveOriginalError && fileStreamError) { + throw fileStreamError; + } + return; + } + + let closeError: Error | null = null; + await new Promise((resolve) => { + const onClose = () => { + cleanup(); + resolve(); + }; + const onCloseError = (error: Error) => { + closeError ??= error; + }; + const cleanup = () => { + fileStream.off("close", onClose); + fileStream.off("error", onCloseError); + }; + + fileStream.once("close", onClose); + fileStream.on("error", onCloseError); + closeAction(); + }); + + const streamError = closeError ?? fileStreamError; + if (!preserveOriginalError && streamError) { + throw streamError; + } + }; const reader = response.body?.getReader(); if (!reader) { - fileStream.destroy(); + await awaitFileStreamClose(() => fileStream.destroy(), true); throw new Error(`Failed to read ${artifact.filename}`); } - let writeCompleted = false; try { while (true) { if (fileStreamError) { @@ -466,7 +500,7 @@ class ModelService extends EventEmitter { if (done) break; if (abortController.signal.aborted) { - fileStream.close(); + await awaitFileStreamClose(() => fileStream.close(), true); if (fs.existsSync(artifactPath)) { fs.unlinkSync(artifactPath); } @@ -513,26 +547,14 @@ class ModelService extends EventEmitter { } } - await new Promise((resolve, reject) => { - const onFinish = () => { - fileStream.off("error", onFinishError); - resolve(); - }; - const onFinishError = (error: Error) => { - fileStream.off("finish", onFinish); - reject(error); - }; - - fileStream.once("finish", onFinish); - fileStream.once("error", onFinishError); - fileStream.end(); - }); - writeCompleted = true; + await awaitFileStreamClose(() => fileStream.end()); + } catch (error) { + if (!fileStream.closed && !fileStream.destroyed) { + await awaitFileStreamClose(() => fileStream.destroy(), true); + } + throw error; } finally { fileStream.off("error", onFileStreamError); - if (!writeCompleted && !fileStream.closed) { - fileStream.destroy(); - } } if (artifact.checksum) {