From 1d3b919f703799634cb30bb2a2123cda230cae64 Mon Sep 17 00:00:00 2001 From: James Rowdy Date: Fri, 7 Aug 2026 06:28:45 +0800 Subject: [PATCH 1/2] Add custom vocabulary prompt for Whisper models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Words the transcriber gets wrong (names, jargon, product terms) can now be listed in the Transforms tab under "Custom vocabulary". The terms are tokenized and injected via DecodingOptions.promptTokens — WhisperKit's equivalent of Whisper's initial_prompt/hotwords — so the decoder treats them as previously transcribed context and is measurably more likely to reproduce their exact spelling and casing. - HexSettings: customVocabulary / customVocabularyEnabled (persisted, backward-compatible defaults via the settings schema) - CustomVocabularyPrompt: parses comma/newline-separated terms and builds a labeled prompt ("Vocabulary: ...") capped at 220 chars, trimming whole terms rather than splitting mid-word (Whisper's decoder context is ~224 tokens shared with prefill, so over-long prompts hurt) - TranscriptionClient.transcribe takes an optional vocabulary prompt and encodes it with the loaded model's tokenizer; a nil tokenizer or empty vocabulary leaves DecodingOptions untouched - No-op for Parakeet, which does not support prompt conditioning; the settings UI notes this when a Parakeet model is selected - Tests cover term parsing, enable/disable, empty input, and prompt trimming (run via the Hex scheme alongside the other HexTests) --- .changeset/custom-vocabulary-prompt.md | 5 ++ Hex/Clients/TranscriptionClient.swift | 16 +++++- .../Remappings/WordRemappingsView.swift | 46 ++++++++++++++++ Hex/Features/Settings/SettingsFeature.swift | 10 ++++ .../Transcription/TranscriptionFeature.swift | 8 ++- .../Models/CustomVocabularyPrompt.swift | 48 +++++++++++++++++ .../HexCore/Settings/HexSettings.swift | 16 ++++-- HexTests/CustomVocabularyPromptTests.swift | 54 +++++++++++++++++++ 8 files changed, 197 insertions(+), 6 deletions(-) create mode 100644 .changeset/custom-vocabulary-prompt.md create mode 100644 HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift create mode 100644 HexTests/CustomVocabularyPromptTests.swift 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..47082bcf6 --- /dev/null +++ b/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift @@ -0,0 +1,48 @@ +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 } + + var prompt = "Vocabulary: \(terms.joined(separator: ", "))." + while prompt.count > maxPromptLength, let lastComma = prompt.lastIndex(of: ",") { + prompt = String(prompt[prompt.startIndex.. { 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..84e401d9d --- /dev/null +++ b/HexTests/CustomVocabularyPromptTests.swift @@ -0,0 +1,54 @@ +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")) + } +} From e80e263d21e421534869af6235c9f951da26c786 Mon Sep 17 00:00:00 2001 From: James Rowdy Date: Fri, 7 Aug 2026 06:46:13 +0800 Subject: [PATCH 2/2] Fix prompt-length cap for over-long single/final terms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CodeRabbit review on #281: the comma-trim loop only enforced maxPromptLength when a comma existed to trim back to, so a single over-long term (or an over-long final term) shipped an over-length prompt — defeating the guard that protects Whisper's ~224-token decoder context. Build the prompt incrementally from whole terms that fit, skipping any term that would exceed the cap and returning nil when none do. Terms are still never split mid-word. Adds tests for the single-overlong, overlong-tail, and overlong-middle cases. --- .../Models/CustomVocabularyPrompt.swift | 20 +++++++++++--- HexTests/CustomVocabularyPromptTests.swift | 27 +++++++++++++++++++ 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift b/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift index 47082bcf6..4409022fc 100644 --- a/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift +++ b/HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift @@ -39,10 +39,22 @@ public enum CustomVocabularyPrompt { let terms = parseTerms(vocabulary) guard !terms.isEmpty else { return nil } - var prompt = "Vocabulary: \(terms.joined(separator: ", "))." - while prompt.count > maxPromptLength, let lastComma = prompt.lastIndex(of: ",") { - prompt = String(prompt[prompt.startIndex..