feat: WhisperKit transcriber biases decoding with vocabulary prompt tokens

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-10-01 15:03:04 +02:00
parent c8e16145e1
commit af1acdc0fa
4 changed files with 69 additions and 2 deletions
@@ -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)
+    }
+}