33 lines
1.3 KiB
Swift
33 lines
1.3 KiB
Swift
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)
|
|
}
|
|
}
|