Add custom vocabulary prompt for Whisper models - #281
Open
jrxwdy wants to merge 2 commits into
Open
Conversation
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)
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThe change adds persisted custom vocabulary settings, a vocabulary editor, prompt construction and truncation, and WhisperKit integration that converts the prompt into decoder tokens. Parakeet models display a notice and do not use the prompt. ChangesCustom vocabulary support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant WordRemappingsView
participant SettingsFeature
participant TranscriptionFeature
participant TranscriptionClientLive
participant WhisperKit
WordRemappingsView->>SettingsFeature: Update vocabulary settings
SettingsFeature->>TranscriptionFeature: Persist custom vocabulary
TranscriptionFeature->>TranscriptionClientLive: Pass custom vocabulary prompt
TranscriptionClientLive->>WhisperKit: Encode prompt into promptTokens
WhisperKit-->>TranscriptionClientLive: Transcribe with decoding options
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.changeset/custom-vocabulary-prompt.md:
- Line 5: Update the changeset entry to append the relevant GitHub issue or pull
request number in `(`#123`)` format, using the actual reference associated with
the custom vocabulary prompt change.
In `@Hex/Features/Remappings/WordRemappingsView.swift`:
- Around line 80-105: Update the custom vocabulary controls in the settings view
so both the “Enabled” Toggle and vocabulary TextField are disabled when
ParakeetModel(rawValue: store.hexSettings.selectedModel) is non-nil. Preserve
the existing customVocabularyEnabled condition for the text field while
combining it with the Parakeet check.
In `@HexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swift`:
- Around line 42-45: Update the prompt-building logic in CustomVocabularyPrompt
to enforce maxPromptLength even when the vocabulary has a single overlong term
or an overlong final term. Construct the prompt incrementally from complete
terms that fit within the limit, and return nil when no term can be included;
preserve the existing punctuation and formatting for included terms.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 613909f3-e670-4b3d-8965-358c9dff9c3d
📒 Files selected for processing (8)
.changeset/custom-vocabulary-prompt.mdHex/Clients/TranscriptionClient.swiftHex/Features/Remappings/WordRemappingsView.swiftHex/Features/Settings/SettingsFeature.swiftHex/Features/Transcription/TranscriptionFeature.swiftHexCore/Sources/HexCore/Models/CustomVocabularyPrompt.swiftHexCore/Sources/HexCore/Settings/HexSettings.swiftHexTests/CustomVocabularyPromptTests.swift
| "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. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add the GitHub issue or PR reference.
Append the actual issue or PR number in (#123) format.
Proposed fix
-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.
+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. (`#123`)As per coding guidelines, include a GitHub issue or PR number in (#123) format.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In @.changeset/custom-vocabulary-prompt.md at line 5, Update the changeset entry
to append the relevant GitHub issue or pull request number in `(`#123`)` format,
using the actual reference associated with the custom vocabulary prompt change.
Source: Coding guidelines
Comment on lines
+80
to
+105
| 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() |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Disable custom vocabulary controls for Parakeet.
The notice does not prevent edits when the selected model is Parakeet. TranscriptionClient ignores the setting on that path. Disable the toggle and text field when ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil.
Proposed fix
.toggleStyle(.switch)
.controlSize(.small)
+ .disabled(ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil)
}
@@
- .disabled(!store.hexSettings.customVocabularyEnabled)
+ .disabled(
+ !store.hexSettings.customVocabularyEnabled ||
+ ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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() | |
| Toggle( | |
| "Enabled", | |
| isOn: Binding( | |
| get: { store.hexSettings.customVocabularyEnabled }, | |
| set: { store.send(.setCustomVocabularyEnabled($0)) } | |
| ) | |
| ) | |
| .toggleStyle(.switch) | |
| .controlSize(.small) | |
| .disabled(ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil) | |
| } | |
| 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 || | |
| ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil | |
| ) | |
| if ParakeetModel(rawValue: store.hexSettings.selectedModel) != nil { | |
| Text("Custom vocabulary only applies to Whisper models. Your selected model is Parakeet.") | |
| .settingsCaption() |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@Hex/Features/Remappings/WordRemappingsView.swift` around lines 80 - 105,
Update the custom vocabulary controls in the settings view so both the “Enabled”
Toggle and vocabulary TextField are disabled when ParakeetModel(rawValue:
store.hexSettings.selectedModel) is non-nil. Preserve the existing
customVocabularyEnabled condition for the text field while combining it with the
Parakeet check.
CodeRabbit review on kitlangton#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.
Author
|
Addressed CodeRabbit's finding on the prompt-length cap (e80e263): the comma-trim loop only enforced |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Hex gets names and jargon wrong because the Whisper decoder has no way to know the user's domain vocabulary. Whisper's canonical fix for this is the
initial_prompt/ hotwords mechanism, and WhisperKit already exposes it viaDecodingOptions.promptTokens— Hex just never wires anything into it.This PR adds a Custom vocabulary section to the Transforms tab: a comma/newline-separated list of terms (names, jargon, product words) that Hex tokenizes with the loaded model's tokenizer and injects as
promptTokensbefore transcription. The decoder treats them as previously transcribed context and is measurably more likely to reproduce their exact spelling/casing.CustomVocabularyPrompt(HexCore): parses terms, builds a labeled prompt (Vocabulary: Langton, TCA.), caps it at 220 chars trimming whole terms — Whisper's decoder context is ~224 tokens shared with prefill, so over-long prompts degrade rather than helpTranscriptionClient.transcribegains an optional vocabulary prompt parameter; nil prompt or nil tokenizer leavesDecodingOptionsuntouched (zero behavior change for existing users)Test plan
HexCorebuilds clean via SwiftPM (swift build)HexTests/CustomVocabularyPromptTests.swift(XCTest) alongside the existing app-level tests — runs withxcodebuild test -scheme Hex(needs full Xcode; my machine only has Command Line Tools, so I couldn't run the app scheme locally)Happy to open an issue first per the contributing note — figured a small, self-contained fix with tests was easier to evaluate as a PR. No worries if you'd rather take it as an issue report.
Summary by CodeRabbit