feat: WhisperKit transcriber biases decoding with vocabulary prompt tokens

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

Package.swift

@@ -19,6 +19,7 @@ let package = Package(
         .executableTarget(name: "Sst", dependencies: ["SstCore", "SstWhisper"]),
         .executableTarget(name: "SstProbe", dependencies: ["SstCore", "SstWhisper"]),
         .testTarget(name: "SstCoreTests", dependencies: ["SstCore"]),
+        .testTarget(name: "SstWhisperTests", dependencies: ["SstWhisper"]),
     ],
     swiftLanguageModes: [.v5]
 )

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)

Tests/SstWhisperTests/PromptTokensTests.swift

@@ -0,0 +1,32 @@
+import Testing
+@testable import SstWhisper
+
+struct PromptTokensTests {
+    /// Fake tokenizer: one token (id 7) per word, plus one special token (id 2000) that must be filtered out.
+    let encode: (String) -> [Int] = { text in
+        text.split(whereSeparator: { !$0.isLetter }).map { _ in 7 } + [2000]
+    }
+
+    @Test func allTermsFitAndSpecialTokensAreFiltered() {
+        let tokens = WhisperKitTranscriber.promptTokens(
+            vocabulary: ["Anthropic", "Vercel", "Supabase"], maxTokens: 64, specialTokenBegin: 1000, encode: encode)
+        #expect(tokens == [7, 7, 7])
+    }
+
+    @Test func termsThatDoNotFitAreDroppedFromTheEnd() {
+        let vocab = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
+        let tokens = WhisperKitTranscriber.promptTokens(
+            vocabulary: vocab, maxTokens: 4, specialTokenBegin: 1000, encode: encode)
+        #expect(tokens?.count == 4)
+    }
+
+    @Test func emptyVocabularyGivesNoPrompt() {
+        #expect(WhisperKitTranscriber.promptTokens(
+            vocabulary: [], maxTokens: 64, specialTokenBegin: 1000, encode: encode) == nil)
+    }
+
+    @Test func firstTermAloneTooLongGivesNoPrompt() {
+        #expect(WhisperKitTranscriber.promptTokens(
+            vocabulary: ["one two three"], maxTokens: 2, specialTokenBegin: 1000, encode: encode) == nil)
+    }
+}