feat: WhisperKit transcriber biases decoding with vocabulary prompt tokens

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-10-01 15:03:04 +02:00
co-authored by Claude Sonnet 5.5
parent c8e16145e1
commit af1acdc0fa
4 changed files with 69 additions and 2 deletions
+32 -1
View File
@@ -77,15 +77,46 @@ public final class WhisperKitTranscriber: Transcriber {
}
public func transcribe(_ samples: [Float]) async throws -> String {
try await transcribe(samples, vocabulary: [])
}
/// Terms are added one by one while the encoded prompt stays within `maxTokens`.
/// Returns nil when nothing fits or the list is empty.
static func promptTokens(
vocabulary: [String],
maxTokens: Int = 64,
specialTokenBegin: Int,
encode: (String) -> [Int]
) -> [Int]? {
var kept: [String] = []
var tokens: [Int] = []
for term in vocabulary {
let candidate = kept + [term]
let encoded = encode(" " + candidate.joined(separator: ", ")).filter { $0 < specialTokenBegin }
if encoded.count > maxTokens { break }
kept = candidate
tokens = encoded
}
return tokens.isEmpty ? nil : tokens
}
public func transcribe(_ samples: [Float], vocabulary: [String]) async throws -> String {
guard let pipeline else { throw TranscriberError.notLoaded }
// language nil + detectLanguage: Whisper picks NL or EN per utterance.
let options = DecodingOptions(
var options = DecodingOptions(
task: .transcribe,
language: nil,
detectLanguage: true,
skipSpecialTokens: true,
withoutTimestamps: true
)
if !vocabulary.isEmpty, let tokenizer = pipeline.tokenizer {
options.promptTokens = Self.promptTokens(
vocabulary: vocabulary,
specialTokenBegin: tokenizer.specialTokens.specialTokenBegin,
encode: { tokenizer.encode(text: $0) }
)
}
let results = try await pipeline.transcribe(audioArray: samples, decodeOptions: options)
return results
.map(\.text)