From af1acdc0fa34167edab0cf72bf9bd6a2c54d2727 Mon Sep 17 00:00:00 2001
From: Theo Weijmer <t.weijmer@pensare.nl>
Date: Thu, 1 Oct 2026 15:03:04 +0200
Subject: [PATCH] feat: WhisperKit transcriber biases decoding with vocabulary
 prompt tokens

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
---
 Package.swift                                 |  1 +
 Sources/SstProbe/TranscribeProbe.swift        |  5 ++-
 .../SstWhisper/WhisperKitTranscriber.swift    | 33 ++++++++++++++++++-
 Tests/SstWhisperTests/PromptTokensTests.swift | 32 ++++++++++++++++++
 4 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 Tests/SstWhisperTests/PromptTokensTests.swift

diff --git a/Package.swift b/Package.swift
index 127d1c3..c8c391c 100644
--- a/Package.swift
+++ b/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]
 )
diff --git a/Sources/SstProbe/TranscribeProbe.swift b/Sources/SstProbe/TranscribeProbe.swift
index f678751..914862c 100644
--- a/Sources/SstProbe/TranscribeProbe.swift
+++ b/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 {
diff --git a/Sources/SstWhisper/WhisperKitTranscriber.swift b/Sources/SstWhisper/WhisperKitTranscriber.swift
index 6bf4916..b64781f 100644
--- a/Sources/SstWhisper/WhisperKitTranscriber.swift
+++ b/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)
diff --git a/Tests/SstWhisperTests/PromptTokensTests.swift b/Tests/SstWhisperTests/PromptTokensTests.swift
new file mode 100644
index 0000000..1ac2169
--- /dev/null
+++ b/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)
+    }
+}