Add custom vocabulary prompt for Whisper models

.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.

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)

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")

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

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

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
+	}
+}

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<Value: Codable & Sendable> {
 	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()
 	]
 }

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.")
+	}
+}