feat: WhisperKit transcriber biases decoding with vocabulary prompt tokens

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

Sources/SstProbe/TranscribeProbe.swift

@@ -25,8 +25,11 @@ func runTranscribeProbe() async {
     }
     print("RMS per seconde:", perSecond.joined(separator: " "))
     do {
+        let vocabulary = (ProcessInfo.processInfo.environment["SST_VOCAB"] ?? "")
+            .split(separator: ",").map { $0.trimmingCharacters(in: .whitespaces) }.filter { !$0.isEmpty }
+        if !vocabulary.isEmpty { print("voorrang: \(vocabulary.joined(separator: ", "))") }
         let started = Date.now
-        let text = try await transcriber.transcribe(samples)
+        let text = try await transcriber.transcribe(samples, vocabulary: vocabulary)
         let elapsed = Date.now.timeIntervalSince(started)
         print(String(format: "(%.2f s) ", elapsed) + text)
     } catch {

Sources/SstWhisper/WhisperKitTranscriber.swift

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