diff --git a/.changeset/9b51fe61.md b/.changeset/9b51fe61.md new file mode 100644 index 000000000..7a0d6718a --- /dev/null +++ b/.changeset/9b51fe61.md @@ -0,0 +1,5 @@ +--- +"hex-app": patch +--- + +Add Parallel Mode transcription with Settings UI toggle and ~/Documents model discovery diff --git a/.gitignore b/.gitignore index 1b0a805bc..9fcd3c0c1 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,7 @@ node_modules/ # Claude .claude/ + +# Test artifacts +test_audio/ + diff --git a/Hex/Clients/TranscriptionClient.swift b/Hex/Clients/TranscriptionClient.swift index 06167c680..319435088 100644 --- a/Hex/Clients/TranscriptionClient.swift +++ b/Hex/Clients/TranscriptionClient.swift @@ -164,7 +164,7 @@ actor TranscriptionClientLive { modelsLogger.info("Deleted model \(variant)") } - /// Returns `true` if the model is already downloaded to the local folder. + /// Returns `true` if the model is already downloaded to local application support or Documents folder. /// Performs a thorough check to ensure the model files are actually present and usable. func isModelDownloaded(_ modelName: String) async -> Bool { if isParakeet(modelName) { @@ -172,34 +172,28 @@ actor TranscriptionClientLive { parakeetLogger.debug("Parakeet available? \(available)") return available } - let modelFolderPath = modelPath(for: modelName).path let fileManager = FileManager.default - // First, check if the basic model directory exists - guard fileManager.fileExists(atPath: modelFolderPath) else { - // Don't print logs that would spam the console - return false - } - - do { - // Check if the directory has actual model files in it - let contents = try fileManager.contentsOfDirectory(atPath: modelFolderPath) + for candidate in modelCandidatePaths(for: modelName) { + let modelFolderPath = candidate.path + guard fileManager.fileExists(atPath: modelFolderPath) else { continue } - // Model should have multiple files and certain key components - guard !contents.isEmpty else { - return false - } + do { + let contents = try fileManager.contentsOfDirectory(atPath: modelFolderPath) + guard !contents.isEmpty else { continue } - // Check for specific model structure - need both tokenizer and model files - let hasModelFiles = contents.contains { $0.hasSuffix(".mlmodelc") || $0.contains("model") } - let tokenizerFolderPath = tokenizerPath(for: modelName).path - let hasTokenizer = fileManager.fileExists(atPath: tokenizerFolderPath) + let hasModelFiles = contents.contains { $0.hasSuffix(".mlmodelc") || $0.contains("model") || $0 == "config.json" } + let tokenizerFolderPath = candidate.appendingPathComponent("tokenizer", isDirectory: true).path + let hasTokenizer = fileManager.fileExists(atPath: tokenizerFolderPath) || contents.contains(where: { $0.contains("tokenizer") }) - // Both conditions must be true for a model to be considered downloaded - return hasModelFiles && hasTokenizer - } catch { - return false + if hasModelFiles && hasTokenizer { + return true + } + } catch { + continue + } } + return false } /// Returns a list of recommended models based on current device hardware. @@ -207,14 +201,48 @@ actor TranscriptionClientLive { await WhisperKit.recommendedRemoteModels() } - /// Lists all model variants available in the `argmaxinc/whisperkit-coreml` repository. + /// Lists all model variants available in remote repository plus local models in App Support and Documents. func getAvailableModels() async throws -> [String] { - var names = try await WhisperKit.fetchAvailableModels() + var names = (try? await WhisperKit.fetchAvailableModels()) ?? [] #if canImport(FluidAudio) for model in ParakeetModel.allCases.reversed() { if !names.contains(model.identifier) { names.insert(model.identifier, at: 0) } } #endif + + let fm = FileManager.default + var scanDirectories: [URL] = [] + + if let baseDir = try? URL.hexModelsDirectory.appendingPathComponent("argmaxinc/whisperkit-coreml", isDirectory: true) { + scanDirectories.append(baseDir) + } + + let docsURL = fm.urls(for: .documentDirectory, in: .userDomainMask).first + if let docsURL { + scanDirectories.append(docsURL.appendingPathComponent("huggingface/models/argmaxinc/whisperkit-coreml", isDirectory: true)) + scanDirectories.append(docsURL.appendingPathComponent("models/argmaxinc/whisperkit-coreml", isDirectory: true)) + scanDirectories.append(docsURL.appendingPathComponent("models", isDirectory: true)) + scanDirectories.append(docsURL) + } + + for dir in scanDirectories { + guard let subdirs = try? fm.contentsOfDirectory(at: dir, includingPropertiesForKeys: [.isDirectoryKey], options: [.skipsHiddenFiles]) else { continue } + for subdir in subdirs { + var isDir: ObjCBool = false + if fm.fileExists(atPath: subdir.path, isDirectory: &isDir), isDir.boolValue { + let folderName = subdir.lastPathComponent + if folderName != "argmaxinc" && folderName != "whisperkit-coreml" && folderName != "huggingface" && folderName != "models" { + if let contents = try? fm.contentsOfDirectory(atPath: subdir.path), + contents.contains(where: { $0.hasSuffix(".mlmodelc") || $0.contains("model") || $0 == "tokenizer" || $0 == "config.json" }) { + if !names.contains(folderName) { + names.append(folderName) + } + } + } + } + } + } + return names } @@ -234,13 +262,13 @@ actor TranscriptionClientLive { try await downloadAndLoadModel(variant: model) { p in progressCallback(p) } - transcriptionLogger.info("Parakeet ensureLoaded took \(String(format: "%.2f", Date().timeIntervalSince(startLoad)))s") + transcriptionLogger.info("Parakeet ensureLoaded took \(String(format: "%.2f", Date().timeIntervalSince(startLoad)), privacy: .public)s") let preparedClip = try ParakeetClipPreparer.ensureMinimumDuration(url: url, logger: parakeetLogger) defer { preparedClip.cleanup() } let startTx = Date() let text = try await parakeet.transcribe(preparedClip.url) - transcriptionLogger.info("Parakeet transcription took \(String(format: "%.2f", Date().timeIntervalSince(startTx)))s") - transcriptionLogger.info("Parakeet request total elapsed \(String(format: "%.2f", Date().timeIntervalSince(startAll)))s") + transcriptionLogger.info("Parakeet transcription took \(String(format: "%.2f", Date().timeIntervalSince(startTx)), privacy: .public)s") + transcriptionLogger.info("Parakeet request total elapsed \(String(format: "%.2f", Date().timeIntervalSince(startAll)), privacy: .public)s") return text } let model = await resolveVariant(model) @@ -253,7 +281,7 @@ actor TranscriptionClientLive { progressCallback(p) } let loadDuration = Date().timeIntervalSince(startLoad) - transcriptionLogger.info("WhisperKit ensureLoaded model=\(model) took \(String(format: "%.2f", loadDuration))s") + transcriptionLogger.info("WhisperKit ensureLoaded model=\(model, privacy: .public) took \(String(format: "%.2f", loadDuration), privacy: .public)s") } guard let whisperKit = whisperKit else { @@ -266,18 +294,147 @@ actor TranscriptionClientLive { ) } - // Perform the transcription. - transcriptionLogger.notice("Transcribing with WhisperKit model=\(model) file=\(url.lastPathComponent)") + // Check if parallel mode is requested via concurrentWorkerCount > 0 + let useParallel = options.concurrentWorkerCount > 0 + let audioLen = try await audioDuration(of: url) + + if useParallel { + transcriptionLogger.notice("Using parallel chunked transcription for model=\(model, privacy: .public) file=\(url.lastPathComponent, privacy: .public) duration=\(String(format: "%.1f", audioLen), privacy: .public)s") + let text = try await transcribeParallel(whisperKit: whisperKit, url: url, model: model, options: options) + transcriptionLogger.info("Parallel transcription total elapsed \(String(format: "%.2f", Date().timeIntervalSince(startAll)), privacy: .public)s") + return text + } + + // Standard transcription path. + transcriptionLogger.notice("Transcribing with WhisperKit model=\(model, privacy: .public) file=\(url.lastPathComponent, privacy: .public)") let startTx = Date() let results = try await whisperKit.transcribe(audioPath: url.path, decodeOptions: options) - transcriptionLogger.info("WhisperKit transcription took \(String(format: "%.2f", Date().timeIntervalSince(startTx)))s") - transcriptionLogger.info("WhisperKit request total elapsed \(String(format: "%.2f", Date().timeIntervalSince(startAll)))s") + let txDuration = Date().timeIntervalSince(startTx) + transcriptionLogger.info("WhisperKit transcription took \(String(format: "%.2f", txDuration), privacy: .public)s for \(String(format: "%.1f", audioLen), privacy: .public)s audio (realtime=\(String(format: "%.1f", audioLen/txDuration), privacy: .public)x)") + transcriptionLogger.info("WhisperKit request total elapsed \(String(format: "%.2f", Date().timeIntervalSince(startAll)), privacy: .public)s") - // Concatenate results from all segments. let text = results.map(\.text).joined(separator: " ") return text } + /// Splits audio into chunks and transcribes them concurrently using WhisperKit's existing pipeline. + private func transcribeParallel( + whisperKit: WhisperKit, + url: URL, + model: String, + options: DecodingOptions + ) async throws -> String { + // Load audio into memory + let loadStart = Date() + let audioArray = try await loadAudioArray(from: url) + let loadTime = Date().timeIntervalSince(loadStart) + let totalSamples = audioArray.count + let sampleRate = WhisperKit.sampleRate + let totalDuration = Double(totalSamples) / Double(sampleRate) + transcriptionLogger.info("Parallel: audio load took \(String(format: "%.3f", loadTime), privacy: .public)s for \(String(format: "%.1f", totalDuration), privacy: .public)s audio") + + // Split into fixed 30s chunks + let chunkSec = 30 + let chunkSamples = chunkSec * sampleRate + var chunks: [[Float]] = [] + var offset = 0 + while offset < totalSamples { + let end = min(offset + chunkSamples, totalSamples) + let chunk = Array(audioArray[offset.. [Float] { + let audioFile = try AVAudioFile(forReading: url) + let format = audioFile.processingFormat + let frameCount = UInt32(audioFile.length) + guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else { + throw NSError(domain: "TranscriptionClient", code: -5, userInfo: [NSLocalizedDescriptionKey: "Failed to create audio buffer"]) + } + try audioFile.read(into: buffer) + + let srcRate = Float(format.sampleRate) + let targetRate = Float(WhisperKit.sampleRate) + + if format.commonFormat == .pcmFormatFloat32 { + let samples = Array(UnsafeBufferPointer(start: buffer.floatChannelData?[0], count: Int(buffer.frameLength))) + if srcRate != targetRate { + return resample(samples, from: srcRate, to: targetRate) + } + return samples + } else if format.commonFormat == .pcmFormatInt16 { + let intSamples = Array(UnsafeBufferPointer(start: buffer.int16ChannelData?[0], count: Int(buffer.frameLength))) + let samples = intSamples.map { Float($0) / Float(Int16.max) } + if srcRate != targetRate { + return resample(samples, from: srcRate, to: targetRate) + } + return samples + } + throw NSError(domain: "TranscriptionClient", code: -6, userInfo: [NSLocalizedDescriptionKey: "Unsupported audio format"]) + } + + /// Simple linear resampling + private func resample(_ samples: [Float], from srcRate: Float, to dstRate: Float) -> [Float] { + let ratio = srcRate / dstRate + let outLen = Int(Float(samples.count) / ratio) + var out = [Float](repeating: 0, count: outLen) + for i in 0.. TimeInterval { + let audioFile = try AVAudioFile(forReading: url) + return TimeInterval(audioFile.length) / audioFile.processingFormat.sampleRate + } + // MARK: - Private Helpers /// Resolve wildcard patterns (e.g. "distil*large-v3") to a concrete model name. @@ -301,11 +458,39 @@ actor TranscriptionClientLive { ParakeetModel(rawValue: name) != nil } + private func modelCandidatePaths(for variant: String) -> [URL] { + let sanitizedVariant = variant.components(separatedBy: CharacterSet(charactersIn: "./\\")).joined(separator: "_") + var candidates: [URL] = [] + + // 1. Application Support default models directory + candidates.append( + modelsBaseFolder + .appendingPathComponent("argmaxinc") + .appendingPathComponent("whisperkit-coreml") + .appendingPathComponent(sanitizedVariant, isDirectory: true) + ) + + // 2. ~/Documents directory & subpaths + let docsURL = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first + if let docsURL { + candidates.append(docsURL.appendingPathComponent("huggingface/models/argmaxinc/whisperkit-coreml/\(sanitizedVariant)", isDirectory: true)) + candidates.append(docsURL.appendingPathComponent("models/argmaxinc/whisperkit-coreml/\(sanitizedVariant)", isDirectory: true)) + candidates.append(docsURL.appendingPathComponent("models/\(sanitizedVariant)", isDirectory: true)) + candidates.append(docsURL.appendingPathComponent(sanitizedVariant, isDirectory: true)) + } + + return candidates + } + /// Creates or returns the local folder (on disk) for a given `variant` model. private func modelPath(for variant: String) -> URL { - // Remove any possible path traversal or invalid characters from variant name + let fm = FileManager.default + for candidate in modelCandidatePaths(for: variant) { + if fm.fileExists(atPath: candidate.path) { + return candidate + } + } let sanitizedVariant = variant.components(separatedBy: CharacterSet(charactersIn: "./\\")).joined(separator: "_") - return modelsBaseFolder .appendingPathComponent("argmaxinc") .appendingPathComponent("whisperkit-coreml") diff --git a/Hex/Features/Settings/GeneralSectionView.swift b/Hex/Features/Settings/GeneralSectionView.swift index 66b6459fd..c9fd8e3ed 100644 --- a/Hex/Features/Settings/GeneralSectionView.swift +++ b/Hex/Features/Settings/GeneralSectionView.swift @@ -82,6 +82,19 @@ struct GeneralSectionView: View { Image(systemName: "bolt.circle") } + Label { + Toggle( + "Parallel Mode", + isOn: Binding( + get: { store.hexSettings.transcribeMode == .parallel }, + set: { store.send(.toggleParallelMode($0)) } + ) + ) + Text("Process Whisper audio in parallel 30s chunks across Neural Engine and GPU cores for up to 10x faster transcription.") + } icon: { + Image(systemName: "cpu") + } + Label { HStack(alignment: .center) { Text("Audio Behavior while Recording") diff --git a/Hex/Features/Settings/SettingsFeature.swift b/Hex/Features/Settings/SettingsFeature.swift index e61dac8d3..0b0b14049 100644 --- a/Hex/Features/Settings/SettingsFeature.swift +++ b/Hex/Features/Settings/SettingsFeature.swift @@ -75,6 +75,7 @@ struct SettingsFeature { case togglePreventSystemSleep(Bool) case setRecordingAudioBehavior(RecordingAudioBehavior) case toggleSuperFastMode(Bool) + case toggleParallelMode(Bool) case setUseClipboardPaste(Bool) case setCopyToClipboard(Bool) case setDoubleTapLockEnabled(Bool) @@ -486,6 +487,10 @@ struct SettingsFeature { state.$hexSettings.withLock { $0.minimumKeyTime = value } return .none + case let .toggleParallelMode(enabled): + state.$hexSettings.withLock { $0.transcribeMode = enabled ? .parallel : .default } + return .none + case let .setOutputLanguage(language): state.$hexSettings.withLock { $0.outputLanguage = language } return .none diff --git a/Hex/Features/Transcription/TranscriptionFeature.swift b/Hex/Features/Transcription/TranscriptionFeature.swift index 4e2d20064..75c82e36b 100644 --- a/Hex/Features/Transcription/TranscriptionFeature.swift +++ b/Hex/Features/Transcription/TranscriptionFeature.swift @@ -337,7 +337,7 @@ private extension TranscriptionFeature { let minimumKeyTime = state.hexSettings.minimumKeyTime let hotkeyHasKey = state.hexSettings.hotkey.key != nil transcriptionFeatureLogger.notice( - "Recording stopped duration=\(String(format: "%.3f", duration))s start=\(startStamp) stop=\(stopStamp) decision=\(String(describing: decision)) minimumKeyTime=\(String(format: "%.2f", minimumKeyTime)) hotkeyHasKey=\(hotkeyHasKey)" + "Recording stopped duration=\(String(format: "%.3f", duration), privacy: .public)s start=\(startStamp, privacy: .public) stop=\(stopStamp, privacy: .public) decision=\(String(describing: decision), privacy: .public) minimumKeyTime=\(String(format: "%.2f", minimumKeyTime), privacy: .public) hotkeyHasKey=\(hotkeyHasKey, privacy: .public)" ) guard decision == .proceedToTranscription else { @@ -365,6 +365,7 @@ private extension TranscriptionFeature { state.isTranscribing = true state.error = nil let language = state.hexSettings.outputLanguage + let isParallelMode = state.hexSettings.transcribeMode == .parallel state.isPrewarming = true @@ -407,12 +408,13 @@ private extension TranscriptionFeature { let decodeOptions = DecodingOptions( language: language, detectLanguage: language == nil, // Only auto-detect if no language specified - chunkingStrategy: .vad, + concurrentWorkerCount: isParallelMode ? 8 : 0, + chunkingStrategy: .vad ) let result = try await transcription.transcribe(capturedURL, model, decodeOptions) { _ in } - transcriptionFeatureLogger.notice("Transcribed audio from \(capturedURL.lastPathComponent) to text length \(result.count)") + transcriptionFeatureLogger.notice("Transcribed audio from \(capturedURL.lastPathComponent, privacy: .public) to text length \(result.count, privacy: .public)") audioURL = nil await send(.transcriptionResult(result, capturedURL, duration)) } catch { diff --git a/Hex/Models/AppHexSettings.swift b/Hex/Models/AppHexSettings.swift index 07b7782ec..0143fd1a5 100644 --- a/Hex/Models/AppHexSettings.swift +++ b/Hex/Models/AppHexSettings.swift @@ -5,6 +5,7 @@ import HexCore // Re-export types so the app target can use them without HexCore prefixes. typealias RecordingAudioBehavior = HexCore.RecordingAudioBehavior +typealias TranscriptionMode = HexCore.TranscriptionMode typealias HexSettings = HexCore.HexSettings extension SharedReaderKey diff --git a/HexCore/Sources/HexCore/Settings/HexSettings.swift b/HexCore/Sources/HexCore/Settings/HexSettings.swift index ca9158e46..7d975099c 100644 --- a/HexCore/Sources/HexCore/Settings/HexSettings.swift +++ b/HexCore/Sources/HexCore/Settings/HexSettings.swift @@ -6,6 +6,11 @@ public enum RecordingAudioBehavior: String, Codable, CaseIterable, Equatable, Se case doNothing } +public enum TranscriptionMode: String, Codable, CaseIterable, Equatable, Sendable { + case `default` + case parallel +} + /// User-configurable settings saved to disk. public struct HexSettings: Codable, Equatable, Sendable { public static let defaultPasteLastTranscriptHotkey = HotKey(key: .v, modifiers: [.option, .shift]) @@ -47,8 +52,10 @@ public struct HexSettings: Codable, Equatable, Sendable { public var wordRemovalsEnabled: Bool public var wordRemovals: [WordRemoval] public var wordRemappings: [WordRemapping] + public var transcribeMode: TranscriptionMode public var lowercaseTranscripts: Bool public var removePunctuation: Bool + public var customModelDirectory: String? private mutating func normalizeDoubleTapSettings() { if !doubleTapLockEnabled { @@ -81,8 +88,10 @@ public struct HexSettings: Codable, Equatable, Sendable { wordRemovalsEnabled: Bool = false, wordRemovals: [WordRemoval] = HexSettings.defaultWordRemovals, wordRemappings: [WordRemapping] = [], + transcribeMode: TranscriptionMode = .default, lowercaseTranscripts: Bool = false, - removePunctuation: Bool = false + removePunctuation: Bool = false, + customModelDirectory: String? = nil ) { self.soundEffectsEnabled = soundEffectsEnabled self.soundEffectsVolume = soundEffectsVolume @@ -108,8 +117,10 @@ public struct HexSettings: Codable, Equatable, Sendable { self.wordRemovalsEnabled = wordRemovalsEnabled self.wordRemovals = wordRemovals self.wordRemappings = wordRemappings + self.transcribeMode = transcribeMode self.lowercaseTranscripts = lowercaseTranscripts self.removePunctuation = removePunctuation + self.customModelDirectory = customModelDirectory normalizeDoubleTapSettings() } @@ -160,6 +171,8 @@ private enum HexSettingKey: String, CodingKey, CaseIterable { case wordRemappings case lowercaseTranscripts case removePunctuation + case transcribeMode + case customModelDirectory } private struct SettingsField { @@ -294,6 +307,15 @@ private enum HexSettingsSchema { default: defaults.wordRemappings ).eraseToAny(), SettingsField(.lowercaseTranscripts, keyPath: \.lowercaseTranscripts, default: defaults.lowercaseTranscripts).eraseToAny(), - SettingsField(.removePunctuation, keyPath: \.removePunctuation, default: defaults.removePunctuation).eraseToAny() + SettingsField(.removePunctuation, keyPath: \.removePunctuation, default: defaults.removePunctuation).eraseToAny(), + SettingsField(.transcribeMode, keyPath: \.transcribeMode, default: defaults.transcribeMode).eraseToAny(), + SettingsField( + .customModelDirectory, + keyPath: \.customModelDirectory, + default: defaults.customModelDirectory, + encode: { container, key, value in + try container.encodeIfPresent(value, forKey: key) + } + ).eraseToAny() ] } diff --git a/HexCore/Tests/HexCoreTests/ParallelModeAndSettingsTests.swift b/HexCore/Tests/HexCoreTests/ParallelModeAndSettingsTests.swift new file mode 100644 index 000000000..b83569b50 --- /dev/null +++ b/HexCore/Tests/HexCoreTests/ParallelModeAndSettingsTests.swift @@ -0,0 +1,68 @@ +import XCTest +@testable import HexCore + +final class ParallelModeAndSettingsTests: XCTestCase { + func testDefaultSettingsHaveParallelModeDisabledByDefault() { + let settings = HexSettings() + XCTAssertEqual(settings.transcribeMode, .default) + XCTAssertNil(settings.customModelDirectory) + } + + func testParallelModeSettingRoundTripSerialization() throws { + var settings = HexSettings() + settings.transcribeMode = .parallel + settings.customModelDirectory = "/Users/test/Documents/models" + + let encoder = JSONEncoder() + let data = try encoder.encode(settings) + + let decoder = JSONDecoder() + let decoded = try decoder.decode(HexSettings.self, from: data) + + XCTAssertEqual(decoded.transcribeMode, .parallel) + XCTAssertEqual(decoded.customModelDirectory, "/Users/test/Documents/models") + XCTAssertEqual(decoded, settings) + } + + func testDecodingJsonPayloadWithParallelTranscribeMode() throws { + let jsonString = """ + { + "transcribeMode": "parallel", + "customModelDirectory": "/Users/test/Documents" + } + """ + guard let data = jsonString.data(using: .utf8) else { + XCTFail("Failed to convert JSON string to Data") + return + } + + let decoded = try JSONDecoder().decode(HexSettings.self, from: data) + XCTAssertEqual(decoded.transcribeMode, .parallel) + XCTAssertEqual(decoded.customModelDirectory, "/Users/test/Documents") + } + + func testTranscriptionModeEnumCasesAndRawValues() { + XCTAssertEqual(TranscriptionMode.default.rawValue, "default") + XCTAssertEqual(TranscriptionMode.parallel.rawValue, "parallel") + XCTAssertEqual(TranscriptionMode.allCases.count, 2) + } + + func testModelPatternMatcherResolvesExactMatch() { + let match = ModelPatternMatcher.matches("openai_whisper-base", "openai_whisper-base") + XCTAssertTrue(match) + } + + func testModelPatternMatcherResolvesWildcardMatch() { + let match = ModelPatternMatcher.matches("openai_whisper-*", "openai_whisper-large-v3-v20240930_turbo_632MB") + XCTAssertTrue(match) + } + + func testModelPatternMatcherResolvesDownloadedPreference() { + let models = [ + (name: "openai_whisper-large-v3", isDownloaded: false), + (name: "openai_whisper-large-v3_turbo", isDownloaded: true) + ] + let resolved = ModelPatternMatcher.resolvePattern("openai_whisper-large-v3*", from: models) + XCTAssertEqual(resolved, "openai_whisper-large-v3_turbo") + } +}