diff --git a/.changeset/custom-vocabulary-prompt.md b/.changeset/custom-vocabulary-prompt.md new file mode 100644 index 000000000..a00122abf --- /dev/null +++ b/.changeset/custom-vocabulary-prompt.md @@ -0,0 +1,5 @@ +--- +"hex-app": minor +--- + +Add custom vocabulary prompt for Whisper models: a new "Custom vocabulary" section in the Transforms tab lets you list names/jargon the transcriber gets wrong. The terms are injected as decoder prompt tokens (Whisper's `initial_prompt` equivalent) before transcription, biasing WhisperKit toward the exact spelling and casing you entered. Applies to Whisper models only; Parakeet ignores the setting. diff --git a/Hex/Clients/TranscriptionClient.swift b/Hex/Clients/TranscriptionClient.swift index 06167c680..98e54ebd8 100644 --- a/Hex/Clients/TranscriptionClient.swift +++ b/Hex/Clients/TranscriptionClient.swift @@ -22,7 +22,9 @@ private let parakeetLogger = HexLog.parakeet struct TranscriptionClient { /// Transcribes an audio file at the specified `URL` using the named `model`. /// Reports transcription progress via `progressCallback`. - var transcribe: @Sendable (URL, String, DecodingOptions, @escaping (Progress) -> Void) async throws -> String + /// `customVocabularyPrompt`, when non-nil, is tokenized and passed to Whisper via + /// `DecodingOptions.promptTokens` to bias decoding toward the user's terms. + var transcribe: @Sendable (URL, String, DecodingOptions, String?, @escaping (Progress) -> Void) async throws -> String /// Ensures a model is downloaded (if missing) and loaded into memory, reporting progress via `progressCallback`. var downloadModel: @Sendable (String, @escaping (Progress) -> Void) async throws -> Void @@ -44,7 +46,7 @@ extension TranscriptionClient: DependencyKey { static var liveValue: Self { let live = TranscriptionClientLive() return Self( - transcribe: { try await live.transcribe(url: $0, model: $1, options: $2, progressCallback: $3) }, + transcribe: { try await live.transcribe(url: $0, model: $1, options: $2, customVocabularyPrompt: $3, progressCallback: $4) }, downloadModel: { try await live.downloadAndLoadModel(variant: $0, progressCallback: $1) }, deleteModel: { try await live.deleteModel(variant: $0) }, isModelDownloaded: { await live.isModelDownloaded($0) }, @@ -225,6 +227,7 @@ actor TranscriptionClientLive { url: URL, model: String, options: DecodingOptions, + customVocabularyPrompt: String? = nil, progressCallback: @escaping (Progress) -> Void ) async throws -> String { let startAll = Date() @@ -267,6 +270,15 @@ actor TranscriptionClientLive { } // Perform the transcription. + // Inject custom vocabulary as decoder prompt tokens (Whisper's `initial_prompt` + // equivalent): the model treats them as previously transcribed context and is + // more likely to reproduce the exact spelling/casing of the listed terms. + var options = options + if let customVocabularyPrompt, + let tokenizer = whisperKit.tokenizer { + options.promptTokens = tokenizer.encode(text: customVocabularyPrompt) + transcriptionLogger.info("Injected custom vocabulary prompt (\(options.promptTokens?.count ?? 0) tokens)") + } transcriptionLogger.notice("Transcribing with WhisperKit model=\(model) file=\(url.lastPathComponent)") let startTx = Date() let results = try await whisperKit.transcribe(audioPath: url.path, decodeOptions: options) diff --git a/Hex/Features/Remappings/WordRemappingsView.swift b/Hex/Features/Remappings/WordRemappingsView.swift index 68f0c0260..ea7624b13 100644 --- a/Hex/Features/Remappings/WordRemappingsView.swift +++ b/Hex/Features/Remappings/WordRemappingsView.swift @@ -13,6 +13,7 @@ struct WordRemappingsView: View { ScrollView { VStack(alignment: .leading, spacing: 24) { previewSection + customVocabularySection wordRulesSection outputFormattingSection } @@ -64,6 +65,51 @@ struct WordRemappingsView: View { } } + private var customVocabularySection: some View { + VStack(alignment: .leading, spacing: 10) { + Text("Custom vocabulary") + .font(.headline) + + GroupBox { + VStack(alignment: .leading, spacing: 10) { + HStack { + Text("Words the transcriber gets wrong (names, jargon). Whisper models are biased toward these spellings before you speak.") + .settingsCaption() + .fixedSize(horizontal: false, vertical: true) + Spacer() + Toggle( + "Enabled", + isOn: Binding( + get: { store.hexSettings.customVocabularyEnabled }, + set: { store.send(.setCustomVocabularyEnabled($0)) } + ) + ) + .toggleStyle(.switch) + .controlSize(.small) + } + + TextField( + "Comma-separated, e.g. Langton, Kit, TCA, WhisperKit", + text: Binding( + get: { store.hexSettings.customVocabulary }, + set: { store.send(.setCustomVocabulary($0)) } + ), + axis: .vertical + ) + .textFieldStyle(.roundedBorder) + .lineLimit(2...4) + .disabled(!store.hexSettings.customVocabularyEnabled) + + if ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil { + Text("Custom vocabulary only applies to Whisper models. Your selected model is Parakeet.") + .settingsCaption() + } + } + .padding(.vertical, 4) + } + } + } + private var wordRulesSection: some View { VStack(alignment: .leading, spacing: 10) { Text("Word rules") diff --git a/Hex/Features/Settings/SettingsFeature.swift b/Hex/Features/Settings/SettingsFeature.swift index e61dac8d3..6d073301b 100644 --- a/Hex/Features/Settings/SettingsFeature.swift +++ b/Hex/Features/Settings/SettingsFeature.swift @@ -115,6 +115,8 @@ struct SettingsFeature { case setRemappingScratchpadFocused(Bool) case setLowercaseTranscripts(Bool) case setRemovePunctuation(Bool) + case setCustomVocabularyEnabled(Bool) + case setCustomVocabulary(String) } @Dependency(\.keyEventMonitor) var keyEventMonitor @@ -411,6 +413,14 @@ struct SettingsFeature { state.$hexSettings.withLock { $0.removePunctuation = enabled } return .none + case let .setCustomVocabularyEnabled(enabled): + state.$hexSettings.withLock { $0.customVocabularyEnabled = enabled } + return .none + + case let .setCustomVocabulary(vocabulary): + state.$hexSettings.withLock { $0.customVocabulary = vocabulary } + return .none + case .startSettingPasteLastTranscriptHotkey: beginCapture(.pasteLastTranscript, state: &state) return .none diff --git a/Hex/Features/Transcription/TranscriptionFeature.swift b/Hex/Features/Transcription/TranscriptionFeature.swift index 4e2d20064..f5c4d4e7b 100644 --- a/Hex/Features/Transcription/TranscriptionFeature.swift +++ b/Hex/Features/Transcription/TranscriptionFeature.swift @@ -365,6 +365,12 @@ private extension TranscriptionFeature { state.isTranscribing = true state.error = nil let language = state.hexSettings.outputLanguage + // Bias Whisper decoding toward the user's custom vocabulary (names, jargon). + // No-op for Parakeet, which doesn't support prompt conditioning. + let customVocabularyPrompt = CustomVocabularyPrompt.makePromptText( + vocabulary: state.hexSettings.customVocabulary, + isEnabled: state.hexSettings.customVocabularyEnabled + ) state.isPrewarming = true @@ -410,7 +416,7 @@ private extension TranscriptionFeature { chunkingStrategy: .vad, ) - let result = try await transcription.transcribe(capturedURL, model, decodeOptions) { _ in } + let result = try await transcription.transcribe(capturedURL, model, decodeOptions, customVocabularyPrompt) { _ in } transcriptionFeatureLogger.notice("Transcribed audio from \(capturedURL.lastPathComponent) to text length \(result.count)") audioURL = nil diff --git a/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift b/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift new file mode 100644 index 000000000..4409022fc --- /dev/null +++ b/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift @@ -0,0 +1,60 @@ +import Foundation + +/// Builds a decoder prompt from the user's custom vocabulary so Whisper-family +/// models are biased toward the exact spelling/casing of names, jargon, and +/// other words the model commonly gets wrong. +/// +/// The prompt is passed to WhisperKit via `DecodingOptions.promptTokens`, +/// which Whisper prepends to the decoder's prefill tokens. The model treats +/// them as previously transcribed text and is measurably more likely to +/// reproduce those spellings — this is the on-device equivalent of Whisper's +/// `initial_prompt` / "hotwords" feature. +public enum CustomVocabularyPrompt { + /// Maximum characters of vocabulary text to inject. Whisper's decoder context + /// is limited (~224 tokens total, shared with prefill tokens), and an + /// over-long prompt degrades rather than improves accuracy, so the prompt is + /// trimmed to the most recent terms that fit. + public static let maxPromptLength = 220 + + /// Parses a free-form vocabulary list into normalized terms. + /// Accepts comma- or newline-separated entries; trims whitespace, drops empties. + public static func parseTerms(_ vocabulary: String) -> [String] { + vocabulary + .components(separatedBy: CharacterSet(charactersIn: ",\n")) + .map { $0.trimmingCharacters(in: .whitespacesAndNewlines) } + .filter { !$0.isEmpty } + } + + /// Builds the prompt text injected ahead of the decoder's prefill tokens. + /// + /// Returns `nil` when the vocabulary is empty or disabled, in which case the + /// caller should leave `DecodingOptions.promptTokens` unset. + /// + /// A short declarative lead-in ("Vocabulary: ...") is used instead of raw + /// comma soup: Whisper is trained on natural prose, so framing the terms as + /// a labeled list conditions it more reliably and lowers the odds of the + /// model echoing stray terms into unrelated transcripts. + public static func makePromptText(vocabulary: String, isEnabled: Bool) -> String? { + guard isEnabled else { return nil } + let terms = parseTerms(vocabulary) + guard !terms.isEmpty else { return nil } + + // Build incrementally from whole terms that fit. Never splits a term, and + // handles a single over-long term (or over-long final term) by dropping it — + // the previous comma-trim loop missed both cases when no comma was present. + let prefix = "Vocabulary: " + let suffix = "." + var includedTerms: [String] = [] + var length = prefix.count + suffix.count + + for term in terms { + let separatorLength = includedTerms.isEmpty ? 0 : 2 // ", " + guard length + separatorLength + term.count <= maxPromptLength else { continue } + includedTerms.append(term) + length += separatorLength + term.count + } + + guard !includedTerms.isEmpty else { return nil } + return prefix + includedTerms.joined(separator: ", ") + suffix + } +} diff --git a/HexCore/Sources/HexCore/Settings/HexSettings.swift b/HexCore/Sources/HexCore/Settings/HexSettings.swift index ca9158e46..9b6e08cb1 100644 --- a/HexCore/Sources/HexCore/Settings/HexSettings.swift +++ b/HexCore/Sources/HexCore/Settings/HexSettings.swift @@ -49,6 +49,8 @@ public struct HexSettings: Codable, Equatable, Sendable { public var wordRemappings: [WordRemapping] public var lowercaseTranscripts: Bool public var removePunctuation: Bool + public var customVocabularyEnabled: Bool + public var customVocabulary: String private mutating func normalizeDoubleTapSettings() { if !doubleTapLockEnabled { @@ -82,7 +84,9 @@ public struct HexSettings: Codable, Equatable, Sendable { wordRemovals: [WordRemoval] = HexSettings.defaultWordRemovals, wordRemappings: [WordRemapping] = [], lowercaseTranscripts: Bool = false, - removePunctuation: Bool = false + removePunctuation: Bool = false, + customVocabularyEnabled: Bool = true, + customVocabulary: String = "" ) { self.soundEffectsEnabled = soundEffectsEnabled self.soundEffectsVolume = soundEffectsVolume @@ -110,6 +114,8 @@ public struct HexSettings: Codable, Equatable, Sendable { self.wordRemappings = wordRemappings self.lowercaseTranscripts = lowercaseTranscripts self.removePunctuation = removePunctuation + self.customVocabularyEnabled = customVocabularyEnabled + self.customVocabulary = customVocabulary normalizeDoubleTapSettings() } @@ -160,7 +166,9 @@ private enum HexSettingKey: String, CodingKey, CaseIterable { case wordRemappings case lowercaseTranscripts case removePunctuation -} + case customVocabularyEnabled + case customVocabulary + } private struct SettingsField { let key: HexSettingKey @@ -294,6 +302,8 @@ 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(.customVocabularyEnabled, keyPath: \.customVocabularyEnabled, default: defaults.customVocabularyEnabled).eraseToAny(), + SettingsField(.customVocabulary, keyPath: \.customVocabulary, default: defaults.customVocabulary).eraseToAny() ] } diff --git a/HexTests/CustomVocabularyPromptTests.swift b/HexTests/CustomVocabularyPromptTests.swift new file mode 100644 index 000000000..14e6d554a --- /dev/null +++ b/HexTests/CustomVocabularyPromptTests.swift @@ -0,0 +1,81 @@ +import XCTest + +@testable import HexCore + +final class CustomVocabularyPromptTests: XCTestCase { + + // MARK: - parseTerms + + func testParseTermsSplitsOnCommasAndNewlines() { + let terms = CustomVocabularyPrompt.parseTerms("Langton, Kit\nTCA,WhisperKit") + XCTAssertEqual(terms, ["Langton", "Kit", "TCA", "WhisperKit"]) + } + + func testParseTermsTrimsWhitespaceAndDropsEmpties() { + let terms = CustomVocabularyPrompt.parseTerms(" Langton ,, \n , TCA ,\n\n") + XCTAssertEqual(terms, ["Langton", "TCA"]) + } + + func testParseTermsEmptyInput() { + XCTAssertEqual(CustomVocabularyPrompt.parseTerms(""), []) + XCTAssertEqual(CustomVocabularyPrompt.parseTerms(" , \n ,"), []) + } + + // MARK: - makePromptText + + func testMakePromptTextDisabledReturnsNil() { + XCTAssertNil(CustomVocabularyPrompt.makePromptText(vocabulary: "Langton", isEnabled: false)) + } + + func testMakePromptTextEmptyVocabularyReturnsNil() { + XCTAssertNil(CustomVocabularyPrompt.makePromptText(vocabulary: "", isEnabled: true)) + XCTAssertNil(CustomVocabularyPrompt.makePromptText(vocabulary: " , \n", isEnabled: true)) + } + + func testMakePromptTextFormatsTermsAsLabeledList() { + let prompt = CustomVocabularyPrompt.makePromptText(vocabulary: "Langton, TCA", isEnabled: true) + XCTAssertEqual(prompt, "Vocabulary: Langton, TCA.") + } + + func testMakePromptTextTrimsOverlongPromptToWholeTerms() { + // ~30 terms of ~16 chars each exceeds maxPromptLength and must be trimmed + let terms = (1...30).map { "VocabularyTerm\($0)" } + let prompt = CustomVocabularyPrompt.makePromptText( + vocabulary: terms.joined(separator: ", "), + isEnabled: true + ) + guard let prompt else { return XCTFail("expected non-nil prompt") } + XCTAssertLessThanOrEqual(prompt.count, CustomVocabularyPrompt.maxPromptLength) + XCTAssertTrue(prompt.hasPrefix("Vocabulary: ")) + XCTAssertTrue(prompt.hasSuffix(".")) + // Trimming must not cut a term in half: the tail term is removed whole + XCTAssertFalse(prompt.contains("VocabularyTerm30")) + } + + func testMakePromptTextSingleOverlongTermReturnsNil() { + // One term longer than the whole budget: no comma to trim to, and the + // term must never be split — so there is nothing usable to prompt with + let hugeTerm = String(repeating: "a", count: CustomVocabularyPrompt.maxPromptLength) + XCTAssertNil(CustomVocabularyPrompt.makePromptText(vocabulary: hugeTerm, isEnabled: true)) + } + + func testMakePromptTextOverlongFinalTermIsDropped() { + // Short terms fit; a trailing over-long term is dropped rather than + // pushing the prompt over the cap + let hugeTail = String(repeating: "z", count: CustomVocabularyPrompt.maxPromptLength) + let prompt = CustomVocabularyPrompt.makePromptText( + vocabulary: "Langton, TCA, \(hugeTail)", + isEnabled: true + ) + XCTAssertEqual(prompt, "Vocabulary: Langton, TCA.") + } + + func testMakePromptTextSkipsOverlongMiddleTermKeepsRest() { + let hugeMiddle = String(repeating: "m", count: CustomVocabularyPrompt.maxPromptLength) + let prompt = CustomVocabularyPrompt.makePromptText( + vocabulary: "Langton, \(hugeMiddle), TCA", + isEnabled: true + ) + XCTAssertEqual(prompt, "Vocabulary: Langton, TCA.") + } +}