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 d3a53495..45b1dafb 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,76 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 4.5, accuracy: 4.5, setup: "cloud", + runtime: "cloud", provider: "Amical Cloud", providerIcon: "/assets/icon_logo.svg", }, + { + 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", @@ -177,6 +252,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 5.0, accuracy: 2.5, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, @@ -209,6 +285,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 4.0, accuracy: 3.0, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, @@ -242,6 +319,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 3.0, accuracy: 3.8, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, @@ -274,6 +352,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 2.0, accuracy: 4.3, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, @@ -306,6 +385,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 1.5, accuracy: 4.7, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, @@ -338,6 +418,7 @@ export const AVAILABLE_MODELS: AvailableWhisperModel[] = [ speed: 3.5, accuracy: 4.2, setup: "offline", + runtime: "whisper-local", provider: "Local", providerIcon: "/icons/models/local.svg", }, diff --git a/apps/desktop/src/db/models.ts b/apps/desktop/src/db/models.ts index 0738bfd5..94af100a 100644 --- a/apps/desktop/src/db/models.ts +++ b/apps/desktop/src/db/models.ts @@ -224,8 +224,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, @@ -238,6 +238,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"); @@ -251,41 +254,99 @@ 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) + 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 filePath = path.join(modelsDirectory, model.filename); - const fileExists = modelFiles.has(model.filename); + const resolvedFiles = resolveRequiredLocalFiles(model); + const 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++; } @@ -304,10 +365,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/main/managers/recording-manager.ts b/apps/desktop/src/main/managers/recording-manager.ts index e1e4bd19..0d5d31b1 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"; @@ -75,6 +76,7 @@ export class RecordingManager extends EventEmitter { // System audio state tracking private systemAudioMuted: boolean = false; + private transcriptionServiceUnavailableNotified: boolean = false; // Sound muting for current session private soundsMuted: boolean = false; @@ -272,6 +274,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 = []; @@ -302,10 +305,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 @@ -382,10 +385,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 }); } @@ -442,14 +445,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); } @@ -474,10 +477,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, @@ -553,10 +559,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 }); } @@ -573,9 +579,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, @@ -814,12 +825,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 }); } @@ -857,10 +868,32 @@ export class RecordingManager extends EventEmitter { this.audioChunks = []; this.terminationCode = null; this.systemAudioMuted = false; + this.transcriptionServiceUnavailableNotified = false; this.soundsMuted = 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..ba89ac59 --- /dev/null +++ b/apps/desktop/src/pipeline/providers/transcription/parakeet-provider.ts @@ -0,0 +1,961 @@ +import * as ort from "onnxruntime-node"; +import * as path from "node:path"; +import { promises as fs } from "node:fs"; +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, + decodeParakeetTokens, + loadParakeetVocabulary, + ParakeetVocabulary, + ParakeetFeatures, +} from "../../utils/parakeet-feature-extractor"; + +interface ResolvedParakeetPaths { + encoderModelPath: string; + decoderJointModelPath: string; + nemoPreprocessorPath?: 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 tdtPreprocessorSession: ort.InferenceSession | null = null; + private tdtEncoderSession: ort.InferenceSession | null = null; + private tdtDecoderJointSession: ort.InferenceSession | null = null; + + private currentModelId: string | null = null; + private vocabulary: ParakeetVocabulary | null = null; + + private frameBuffer: Float32Array[] = []; + private frameBufferSpeechProbabilities: number[] = []; + private currentSilenceFrameCount = 0; + + private featureSize = 80; + private maxTokensPerStep = 10; + private featureExtractor = new ParakeetFeatureExtractor(80); + + 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 { + await this.releaseSessions(); + this.currentModelId = null; + this.vocabulary = null; + this.reset(); + } + + private async doTranscription(_context: TranscribeContext): Promise { + try { + if (!this.vocabulary) { + throw new AppError( + "Parakeet model is not initialized", + ErrorCodes.WORKER_INITIALIZATION_FAILED, + ); + } + + const bufferedFrames = [...this.frameBuffer]; + const vadProbs = [...this.frameBufferSpeechProbabilities]; + const rawAudio = this.aggregateFrames(bufferedFrames); + + const { audio: speechAudio, segments } = extractSpeechFromVad( + rawAudio, + vadProbs, + ); + + if (speechAudio.length === 0) { + logger.transcription.debug( + "Skipping Parakeet transcription - no speech detected by VAD filter", + ); + this.reset(); + return ""; + } + + logger.transcription.debug("Parakeet VAD filtered audio", { + before: rawAudio.length, + after: speechAudio.length, + segments: segments.length, + }); + + const features = await this.extractFeatures(speechAudio); + const transcript = await this.transcribeTdt(features); + this.reset(); + return transcript; + } 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 extractFeatures( + audioData: Float32Array, + ): Promise { + if (!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 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.vocabulary && + this.currentModelId === requestedId && + this.tdtEncoderSession && + this.tdtDecoderJointSession + ) { + return; + } + + const resolved = await this.resolveModelPaths(requestedId); + const config = await this.loadModelConfig(resolved.configPath); + const newFeatureSize = + typeof config.features_size === "number" ? config.features_size : 128; + const newMaxTokensPerStep = + typeof config.max_tokens_per_step === "number" + ? config.max_tokens_per_step + : 10; + const newFeatureExtractor = new ParakeetFeatureExtractor(newFeatureSize); + + const newVocabulary = await loadParakeetVocabulary(resolved.vocabPath); + + const preferredProviders = + process.platform === "win32" + ? (["dml", "cpu"] as const) + : process.platform === "darwin" + ? (["coreml", "cpu"] as const) + : (["cpu"] as const); + + let preprocessorProviders: readonly string[] | null = null; + 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, + ); + newEncoderSession = encoderResult.session; + encoderProviders = encoderResult.providersUsed; + + 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, + nemoPreprocessorPath: resolved.nemoPreprocessorPath || null, + encoderModelPath: resolved.encoderModelPath, + decoderJointModelPath: resolved.decoderJointModelPath, + vocabPath: resolved.vocabPath, + executionProviders: { + preprocessor: preprocessorProviders, + encoder: encoderProviders, + decoder: decoderProviders, + }, + featureSize: this.featureSize, + maxTokensPerStep: this.maxTokensPerStep, + }); + } + + private async createSessionWithFallback( + modelPath: string, + preferredProviders: readonly string[], + ): Promise<{ + session: ort.InferenceSession; + providersUsed: readonly string[]; + }> { + try { + 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), + }, + ); + const session = await ort.InferenceSession.create(modelPath, { + executionProviders: ["cpu"], + }); + return { session, providersUsed: ["cpu"] }; + } + throw error; + } + } + + private async releaseSessions(): Promise { + if (this.tdtPreprocessorSession) { + await this.tdtPreprocessorSession.release(); + this.tdtPreprocessorSession = null; + } + if (this.tdtEncoderSession) { + await this.tdtEncoderSession.release(); + this.tdtEncoderSession = null; + } + if (this.tdtDecoderJointSession) { + await this.tdtDecoderJointSession.release(); + this.tdtDecoderJointSession = null; + } + } + + private async releaseSession( + session: ort.InferenceSession | null, + ): Promise { + if (session) { + await session.release(); + } + } + + 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 { + 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 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( + (value): value is string => typeof value === "string", + ) + : [downloaded.localPath]; + + const findFile = (pattern: RegExp): string | undefined => + localFiles.find((filePath) => pattern.test(path.basename(filePath))); + + const vocabPath = + findFile(/^vocab\.txt$/i) || path.join(modelDir, "vocab.txt"); + + 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 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 { + encoderModelPath, + decoderJointModelPath, + nemoPreprocessorPath, + vocabPath, + configPath, + }; + } + + 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 {}; + } + + 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 bestIndex; + } + + 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( + 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 frames) { + 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..a59e072b --- /dev/null +++ b/apps/desktop/src/pipeline/utils/parakeet-feature-extractor.ts @@ -0,0 +1,322 @@ +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 DEFAULT_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(numMels: number): Float32Array[] { + const numBins = Math.floor(N_FFT / 2) + 1; + const fbanks: Float32Array[] = Array.from( + { length: numMels }, + () => new Float32Array(numBins), + ); + + const minMel = hzToMel(F_MIN); + const maxMel = hzToMel(F_MAX); + const melPoints = new Float64Array(numMels + 2); + for (let i = 0; i < melPoints.length; i++) { + melPoints[i] = minMel + ((maxMel - minMel) * i) / (numMels + 1); + } + + 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 <= numMels; 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 nMels: number; + private readonly window = buildCenteredHannWindow(); + 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) { + 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 * this.nMels); + 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 < 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 * this.nMels + m] = Math.log(energy + LOG_ZERO_GUARD); + } + } + + const validFrames = Math.min(featuresLength, frameCount); + const normalized = new Float32Array(this.nMels * frameCount); + + for (let m = 0; m < this.nMels; m++) { + let mean = 0; + for (let f = 0; f < validFrames; f++) { + mean += logMel[f * this.nMels + m]; + } + mean /= validFrames; + + let variance = 0; + for (let f = 0; f < validFrames; f++) { + const delta = logMel[f * this.nMels + 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 * this.nMels + m] - mean) * invStd : 0; + } + } + + return { + inputFeatures: normalized, + 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 { + 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); + } + + 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, + }; +} diff --git a/apps/desktop/src/services/model-service.ts b/apps/desktop/src/services/model-service.ts index e223ce3e..22c82276 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, syncModelsForProviderInstance, @@ -155,10 +154,22 @@ 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_MODEL_ID, + "whisper-large-v3", + "whisper-medium", + "whisper-small", + "whisper-base", + "whisper-tiny", + ]; constructor(settingsService: SettingsService) { super(); @@ -204,8 +215,10 @@ 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) => ({ + // 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, @@ -214,6 +227,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( @@ -247,22 +263,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, @@ -297,29 +299,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, + }); } } @@ -357,22 +348,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, @@ -418,6 +395,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; @@ -435,17 +421,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 @@ -465,6 +458,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}`); } @@ -474,14 +471,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, }; @@ -492,85 +512,175 @@ 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(), - }, - }); + let bytesDownloaded = 0; + let lastProgressEmit = 0; + const localFiles: string[] = []; - if (!response.ok) { - throw new Error( - `Failed to download: ${response.status} ${response.statusText}`, - ); - } + for (const artifact of artifacts) { + const artifactPath = path.join(modelDirectory, artifact.filename); - const totalBytes = - parseInt(response.headers.get("content-length") || "0") || model.size; - progress.totalBytes = totalBytes; + const response = await fetch(artifact.downloadUrl, { + signal: abortController.signal, + headers: { + "User-Agent": getUserAgent(), + }, + }); - const fileStream = fs.createWriteStream(downloadPath); - let bytesDownloaded = 0; - let lastProgressEmit = 0; + if (!response.ok) { + throw new Error( + `Failed to download ${artifact.filename}: ${response.status} ${response.statusText}`, + ); + } - const reader = response.body?.getReader(); - if (!reader) { - throw new Error("Failed to get response reader"); - } + const artifactBytes = + parseInt(response.headers.get("content-length") || "0") || + artifact.size || + 0; + if (!artifact.size && artifactBytes > 0) { + progress.totalBytes += artifactBytes; + } - while (true) { - const { done, value } = await reader.read(); + const fileStream = fs.createWriteStream(artifactPath); + let fileStreamError: Error | null = null; + const onFileStreamError = (error: Error) => { + fileStreamError = error; + }; + fileStream.on("error", onFileStreamError); + const awaitFileStreamClose = async ( + closeAction: () => void, + preserveOriginalError = false, + ): Promise => { + if (fileStream.closed) { + if (!preserveOriginalError && fileStreamError) { + throw fileStreamError; + } + return; + } - if (done) break; + 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(); + }); - if (abortController.signal.aborted) { - fileStream.close(); - fs.unlinkSync(downloadPath); - throw new Error("Download cancelled"); + const streamError = closeError ?? fileStreamError; + if (!preserveOriginalError && streamError) { + throw streamError; + } + }; + const reader = response.body?.getReader(); + if (!reader) { + await awaitFileStreamClose(() => fileStream.destroy(), true); + throw new Error(`Failed to read ${artifact.filename}`); } - fileStream.write(value); - bytesDownloaded += value.length; + try { + while (true) { + if (fileStreamError) { + throw fileStreamError; + } + + const { done, value } = await reader.read(); + if (done) break; + + if (abortController.signal.aborted) { + await awaitFileStreamClose(() => fileStream.close(), true); + 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; + } + + 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; + } + } - progress.bytesDownloaded = bytesDownloaded; - progress.progress = Math.round((bytesDownloaded / totalBytes) * 100); + await awaitFileStreamClose(() => fileStream.end()); + } catch (error) { + if (!fileStream.closed && !fileStream.destroyed) { + await awaitFileStreamClose(() => fileStream.destroy(), true); + } + throw error; + } finally { + fileStream.off("error", onFileStreamError); + } - // 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; + 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}`, + ); + } } - } - fileStream.end(); + 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, @@ -587,10 +697,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 @@ -608,14 +727,13 @@ 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 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( @@ -630,9 +748,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)); @@ -683,12 +801,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) @@ -702,30 +844,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) { @@ -743,14 +879,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); }); } @@ -784,12 +925,50 @@ class ModelService extends EventEmitter { return null; } - const normalizedSelection = getSpeechModelSelectionKey(modelId); + const normalizedModelId = + await this.normalizeLegacySpeechModelSelection(modelId); + if (!normalizedModelId) { + return null; + } + + const normalizedSelection = + getSpeechModelSelectionKey(normalizedModelId); if (normalizedSelection !== selection) { await this.settingsService.setDefaultSpeechModel(normalizedSelection); } - return modelId; + return normalizedModelId; + } + + 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( @@ -872,6 +1051,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 @@ -914,7 +1097,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", @@ -1530,40 +1713,43 @@ class ModelService extends EventEmitter { */ async validateAndClearInvalidDefaults(): Promise { // Check default speech model - const defaultSpeechModel = + const storedDefaultSpeechModel = await this.settingsService.getDefaultSpeechModel(); - if (defaultSpeechModel) { + if (storedDefaultSpeechModel) { const speechModelId = - getSpeechModelIdFromStoredSelection(defaultSpeechModel); + getSpeechModelIdFromStoredSelection(storedDefaultSpeechModel); if (!speechModelId) { await this.applySpeechModelSelection(null, "auto-after-deletion", null); } else { - const normalizedSelection = getSpeechModelSelectionKey(speechModelId); - const availableModel = AVAILABLE_MODELS.find( - (m) => m.id === speechModelId, - ); - const isAmicalModel = availableModel?.provider === "Amical Cloud"; - const existsInDb = await modelExists( - getSystemProviderInstanceId(PROVIDER_TYPES.localWhisper), - "speech", - speechModelId, - ); - - if (normalizedSelection !== defaultSpeechModel) { - await this.settingsService.setDefaultSpeechModel(normalizedSelection); - } + const normalizedSpeechModelId = + await this.normalizeLegacySpeechModelSelection(speechModelId); - // Amical cloud models are always valid; local models must exist in DB - if (!isAmicalModel && !existsInDb) { - logger.main.info("Clearing invalid default speech model", { - modelId: speechModelId, - }); - await this.applySpeechModelSelection( - null, - "auto-after-deletion", - speechModelId, + if (normalizedSpeechModelId) { + const normalizedSelection = + getSpeechModelSelectionKey(normalizedSpeechModelId); + const availableModel = AVAILABLE_MODELS.find( + (m) => m.id === normalizedSpeechModelId, ); + const isAmicalModel = availableModel?.setup === "cloud"; + const validDownloadedModels = await this.getValidDownloadedModels(); + const existsLocally = !!validDownloadedModels[normalizedSpeechModelId]; + + if (normalizedSelection !== storedDefaultSpeechModel) { + await this.settingsService.setDefaultSpeechModel(normalizedSelection); + } + + // Amical cloud models are always valid; local models must have a complete bundle. + if (!isAmicalModel && !existsLocally) { + logger.main.info("Clearing invalid default speech model", { + modelId: normalizedSpeechModelId, + }); + await this.applySpeechModelSelection( + null, + "auto-after-deletion", + normalizedSpeechModelId, + ); + } } } } @@ -1630,6 +1816,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/services/transcription-service.ts b/apps/desktop/src/services/transcription-service.ts index f8dcfe4a..e591c339 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 { createRemoteFormattingProvider } from "../pipeline/providers/formatting/remote-formatting-provider-registry"; import type { RemoteFormattingProviderType } from "../pipeline/providers/formatting/remote-formatting-provider-registry"; @@ -43,6 +44,7 @@ import { */ export class TranscriptionService { private whisperProvider: WhisperProvider; + private parakeetProvider: ParakeetProvider; private cloudProvider: AmicalCloudProvider; private currentProvider: TranscriptionProvider | null = null; private streamingSessions = new Map(); @@ -65,6 +67,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; @@ -90,12 +93,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; @@ -107,7 +116,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) { @@ -121,10 +130,20 @@ 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", @@ -179,7 +198,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; } } @@ -213,12 +232,26 @@ 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"); } @@ -319,6 +352,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({ @@ -326,6 +360,7 @@ export class TranscriptionService { speechProbability: speechProbability, context: { sessionId, + modelId: selectedSpeechModelId || undefined, vocabulary: session.context.sharedData.vocabulary, accessibilityContext: session.context.sharedData.accessibilityContext, previousChunk, @@ -428,8 +463,11 @@ 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, @@ -493,11 +531,10 @@ export class TranscriptionService { audioFilePath, hasAudioFile: !!audioFilePath, }); - - const selectedModelId = await this.modelService.getSelectedModel(); + const selectedSpeechModelId = await this.modelService.getSelectedModel(); const speechModelId = usedCloudProvider ? "amical-cloud" - : selectedModelId || "whisper-local"; + : selectedSpeechModelId || "whisper-local"; await createTranscription({ text: completeTranscription, @@ -931,6 +968,7 @@ export class TranscriptionService { audioData: frames[i], speechProbability: vadProbs[i], context: { + modelId: selectedModelId || undefined, sessionId: retrySessionId, vocabulary, language, @@ -949,6 +987,7 @@ export class TranscriptionService { // Flush to get remaining buffered audio const aggregatedTranscription = transcriptionResults.join(""); const finalTranscription = await provider.flush({ + modelId: selectedModelId || undefined, sessionId: retrySessionId, vocabulary, language, @@ -1066,6 +1105,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"); } diff --git a/apps/desktop/src/trpc/routers/models.ts b/apps/desktop/src/trpc/routers/models.ts index d40e000c..f879d72c 100644 --- a/apps/desktop/src/trpc/routers/models.ts +++ b/apps/desktop/src/trpc/routers/models.ts @@ -88,7 +88,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"); @@ -101,6 +101,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, providerType: m.id === "amical-cloud" ? PROVIDER_TYPES.amical 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({}); + }); +});