Skip to content

Commit 02a8edd

Browse files
stikvesstikves
authored
Fix Gemma stop tokens: read additional EOS from tokenizer config
1 parent a43371e commit 02a8edd

5 files changed

Lines changed: 171 additions & 15 deletions

File tree

‎swift/Sources/CoreAILanguageModels/Bundle/LanguageConfig.swift‎

Lines changed: 87 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause
55

66
import CoreAIShared
7+
import Foundation
8+
import Tokenizers
79

810
/// `language` block of `metadata.json` schema 0.2 — LLM-specific config.
911
public struct LanguageConfig: Codable, Sendable, Equatable {
@@ -42,12 +44,96 @@ public struct LanguageConfig: Codable, Sendable, Equatable {
4244
case functionMap = "function_map"
4345
}
4446

45-
public init(from decoder: Decoder) throws {
47+
public init(from decoder: Swift.Decoder) throws {
4648
let c = try decoder.container(keyedBy: CodingKeys.self)
4749
self.tokenizer = try c.decode(String.self, forKey: .tokenizer)
4850
self.vocabSize = try c.decode(Int.self, forKey: .vocabSize)
4951
self.maxContextLength = try c.decode(Int.self, forKey: .maxContextLength)
5052
self.embeddedTokenizer = try c.decodeIfPresent(Bool.self, forKey: .embeddedTokenizer) ?? true
5153
self.functionMap = try c.decodeIfPresent(FunctionMap.self, forKey: .functionMap)
5254
}
55+
56+
// MARK: - Additional Stop Tokens
57+
58+
/// Extract additional stop token IDs from the tokenizer config.
59+
/// Reads `additional_special_tokens` from tokenizer_config.json and
60+
/// cross-references with the tokenizer to get their IDs.
61+
///
62+
/// Also checks for array-valued `eos_token` (some models list multiple).
63+
///
64+
/// Best-effort: returns empty if the file doesn't exist or can't be parsed.
65+
///
66+
/// TODO: Upstream this to swift-transformers as `Tokenizer.additionalEosTokenIds`
67+
/// so we don't need to parse tokenizer_config.json ourselves.
68+
public static func additionalStopTokenIds(
69+
from tokenizerDir: URL,
70+
tokenizer: any Tokenizer
71+
) -> [Int32] {
72+
let configURL = tokenizerDir.appending(path: "tokenizer_config.json")
73+
guard let data = try? Data(contentsOf: configURL),
74+
let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
75+
else {
76+
return []
77+
}
78+
79+
let mainEos = tokenizer.eosTokenId.map { Int32($0) }
80+
var result = Set<Int32>()
81+
82+
// 1. Check additional_special_tokens array
83+
if let specials = json["additional_special_tokens"] as? [Any] {
84+
for item in specials {
85+
// Each item can be a string or a dict with a "content" key
86+
let tokenString: String?
87+
if let s = item as? String {
88+
tokenString = s
89+
} else if let dict = item as? [String: Any],
90+
let content = dict["content"] as? String
91+
{
92+
tokenString = content
93+
} else {
94+
tokenString = nil
95+
}
96+
guard let token = tokenString else { continue }
97+
98+
if let id = tokenizer.convertTokenToId(token) {
99+
let id32 = Int32(id)
100+
if id32 != mainEos {
101+
result.insert(id32)
102+
}
103+
}
104+
}
105+
}
106+
107+
// 2. Check if eos_token is an array (some models list multiple)
108+
if let eosArray = json["eos_token"] as? [String] {
109+
for token in eosArray {
110+
if let id = tokenizer.convertTokenToId(token) {
111+
let id32 = Int32(id)
112+
if id32 != mainEos {
113+
result.insert(id32)
114+
}
115+
}
116+
}
117+
}
118+
119+
// 3. Check added_tokens_decoder for turn-ending special tokens
120+
// (e.g. Gemma's <end_of_turn> ID 106, Qwen's <|im_end|>)
121+
// Only include tokens whose content matches known turn-ending patterns.
122+
let turnEndPatterns = ["end_of_turn", "im_end", "eot_id"]
123+
if let addedTokens = json["added_tokens_decoder"] as? [String: Any] {
124+
for (idString, value) in addedTokens {
125+
guard let dict = value as? [String: Any],
126+
let isSpecial = dict["special"] as? Bool, isSpecial,
127+
let content = dict["content"] as? String,
128+
let id = Int32(idString)
129+
else { continue }
130+
let lower = content.lowercased()
131+
if id != mainEos && turnEndPatterns.contains(where: { lower.contains($0) }) {
132+
result.insert(id)
133+
}
134+
}
135+
}
136+
137+
return Array(result)
138+
}
53139
}

‎swift/Sources/CoreAILanguageModels/DecodingStrategies/DecodingStrategy.swift‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,13 @@ public struct StopSequences: Sendable {
7676
/// Initialize with tokenizer, automatically including EOS tokens
7777
/// - Parameter tokenizer: Tokenizer to extract EOS token from
7878
/// - Parameter additionalSequences: Optional additional stop sequences to include
79-
public init(for tokenizer: any Tokenizer, additionalSequences: [[Int32]] = []) {
79+
/// - Parameter additionalEosTokenIds: Optional additional single-token EOS IDs
80+
/// (e.g. from tokenizer_config.json's `additional_special_tokens`)
81+
public init(
82+
for tokenizer: any Tokenizer,
83+
additionalSequences: [[Int32]] = [],
84+
additionalEosTokenIds: [Int32] = []
85+
) {
8086
var allSequences = additionalSequences
8187

8288
// Collect existing single-token sequences to avoid duplicates
@@ -94,6 +100,14 @@ public struct StopSequences: Sendable {
94100
}
95101
}
96102

103+
// Add additional EOS token IDs (e.g. from tokenizer_config.json)
104+
for token in additionalEosTokenIds {
105+
if !existingTokens.contains(token) {
106+
existingTokens.insert(token)
107+
allSequences.append([token])
108+
}
109+
}
110+
97111
self.sequences = allSequences
98112
self.maxLength = allSequences.map { $0.count }.max() ?? 0
99113
}

‎swift/Sources/CoreAILanguageModels/LanguageModel/CoreAILanguageModel.swift‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ public struct CoreAILanguageModel: LanguageModel {
3737
private let vocabSize: Int?
3838
private let supportsToolCalling: Bool
3939
private let supportsReasoning: Bool
40+
private let additionalEosTokenIds: [Int32]
4041

4142
// MARK: - Protocol Requirements
4243

@@ -56,7 +57,8 @@ public struct CoreAILanguageModel: LanguageModel {
5657
tokenizer: tokenizer,
5758
modelIdentifier: modelIdentifier,
5859
samplingConfig: samplingConfig,
59-
vocabSize: vocabSize
60+
vocabSize: vocabSize,
61+
additionalEosTokenIds: additionalEosTokenIds
6062
)
6163
}
6264

@@ -97,13 +99,15 @@ public struct CoreAILanguageModel: LanguageModel {
9799
tokenizer: any Tokenizer,
98100
modelIdentifier: String = "coreai-model",
99101
samplingConfig: SamplingConfiguration = .greedy,
100-
vocabSize: Int? = nil
102+
vocabSize: Int? = nil,
103+
additionalEosTokenIds: [Int32] = []
101104
) {
102105
self.engine = engine
103106
self.tokenizer = tokenizer
104107
self.modelIdentifier = modelIdentifier
105108
self.samplingConfig = samplingConfig
106109
self.vocabSize = vocabSize
110+
self.additionalEosTokenIds = additionalEosTokenIds
107111
self.supportsToolCalling = CoreAIExecutor.detectToolCallMarkers(using: tokenizer) != nil
108112
self.supportsReasoning =
109113
tokenizer.convertTokenToId("<think>") != nil
@@ -121,6 +125,7 @@ public struct CoreAILanguageModel: LanguageModel {
121125
fileprivate let modelIdentifier: String
122126
fileprivate let samplingConfig: SamplingConfiguration
123127
fileprivate let vocabSize: Int?
128+
fileprivate let additionalEosTokenIds: [Int32]
124129

125130
public static func == (lhs: Configuration, rhs: Configuration) -> Bool {
126131
lhs.modelIdentifier == rhs.modelIdentifier
@@ -140,6 +145,9 @@ public struct CoreAILanguageModel: LanguageModel {
140145
private let modelIdentifier: String
141146
private let samplingConfig: SamplingConfiguration
142147
private let vocabSize: Int?
148+
/// All EOS-like token IDs: the main `eosTokenId` plus any additional
149+
/// stop tokens from tokenizer_config.json (e.g. Gemma's `<end_of_turn>`).
150+
private let eosTokenIds: Set<Int32>
143151
/// Open / close marker pair the model uses for chain-of-thought
144152
/// blocks, discovered from the tokenizer's known token ids at init
145153
/// (see `detectThinkingMarkers`). For models that don't emit
@@ -162,6 +170,14 @@ public struct CoreAILanguageModel: LanguageModel {
162170
self.vocabSize = configuration.vocabSize
163171
self.thinkingMarkers = Self.detectThinkingMarkers(using: configuration.tokenizer)
164172
self.toolCallMarkers = Self.detectToolCallMarkers(using: configuration.tokenizer)
173+
174+
// Build the full set of EOS-like token IDs
175+
var eos = Set<Int32>()
176+
if let id = configuration.tokenizer.eosTokenId {
177+
eos.insert(Int32(id))
178+
}
179+
eos.formUnion(configuration.additionalEosTokenIds)
180+
self.eosTokenIds = eos
165181
}
166182

167183
/// Probes the tokenizer for known reasoning marker pairs. Each
@@ -328,7 +344,8 @@ public struct CoreAILanguageModel: LanguageModel {
328344
inferenceOptions: InferenceOptions(maxTokens: maxTokens)
329345
)
330346

331-
let eosTokenId = tokenizer.eosTokenId
347+
// Use pre-computed set of all EOS-like tokens (main + additional)
348+
let eosTokens = eosTokenIds
332349
// Incremental-decode buffer. After a clean emit, one token is
333350
// retained as context for the next step (see below). During a
334351
// multi-byte sequence that hasn't decoded cleanly yet, multiple
@@ -359,7 +376,7 @@ public struct CoreAILanguageModel: LanguageModel {
359376

360377
for try await output in tokenStream {
361378
let token = output.tokenId
362-
if let eos = eosTokenId, Int(token) == eos {
379+
if eosTokens.contains(token) {
363380
tokenStream.setStopReason(.eos)
364381
break
365382
}
@@ -523,7 +540,10 @@ public struct CoreAILanguageModel: LanguageModel {
523540
}
524541

525542
let strategy = ConstrainedDecodingStrategy(jsonSchema: jsonSchema, vocabSize: vocabSize)
526-
let stopSequences = StopSequences(for: tokenizer)
543+
let stopSequences = StopSequences(
544+
for: tokenizer,
545+
additionalEosTokenIds: Array(eosTokenIds)
546+
)
527547

528548
let stream = try await strategy.decode(
529549
from: .tokens(promptTokens),

‎swift/Sources/CoreAILanguageModels/LanguageModel/CoreAIRunner.swift‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,22 @@ public struct CoreAIRunner {
8080
let tokenizer = try await bundle.loadTokenizer()
8181
tokenizerLoadSpan.end()
8282

83+
// Read additional stop token IDs from tokenizer_config.json
84+
let additionalEos: [Int32]
85+
if let tokenizerDir = bundle.tokenizerPath {
86+
additionalEos = LanguageConfig.additionalStopTokenIds(
87+
from: tokenizerDir, tokenizer: tokenizer)
88+
} else {
89+
additionalEos = []
90+
}
91+
8392
return CoreAILanguageModel(
8493
engine: engine,
8594
tokenizer: tokenizer,
8695
modelIdentifier: bundle.name,
8796
samplingConfig: SamplingConfiguration.greedy,
88-
vocabSize: bundle.vocabSize
97+
vocabSize: bundle.vocabSize,
98+
additionalEosTokenIds: additionalEos
8999
)
90100
}
91101

‎swift/Sources/Tools/llm-runner/LLMRunnerMain.swift‎

Lines changed: 33 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,20 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
361361
"Tokenizer loaded from \(bundle.hasEmbeddedTokenizer ? "embedded bundle" : "HuggingFace")",
362362
component: "Main")
363363

364+
// Read additional stop token IDs from tokenizer_config.json (e.g. <end_of_turn> for Gemma)
365+
let additionalEosTokenIds: [Int32]
366+
if let tokenizerDir = bundle.tokenizerPath {
367+
additionalEosTokenIds = LanguageConfig.additionalStopTokenIds(
368+
from: tokenizerDir, tokenizer: tokenizer)
369+
if !additionalEosTokenIds.isEmpty {
370+
CLILogger.log(
371+
"Found \(additionalEosTokenIds.count) additional stop token(s) from tokenizer config: \(additionalEosTokenIds)",
372+
component: "Main")
373+
}
374+
} else {
375+
additionalEosTokenIds = []
376+
}
377+
364378
CLILogger.log("Model loaded successfully:", component: "Main")
365379
CLILogger.log(" Name: \(modelName)", component: "Main")
366380
CLILogger.log(" Source: model bundle", component: "Main")
@@ -494,7 +508,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
494508
samplingConfiguration: samplingConfiguration,
495509
maxTokens: maxTokens,
496510
actualInputTokens: actualInputTokens,
497-
modelVocabSize: modelVocabSize
511+
modelVocabSize: modelVocabSize,
512+
additionalEosTokenIds: additionalEosTokenIds
498513
)
499514
} else {
500515
// Generate text (timing handled by decoding strategies)
@@ -505,7 +520,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
505520
// Encode stop tokens to sequences
506521
let stopSequences = try validateAndEncodeStopTokens(
507522
stopTokens: stopTokens,
508-
tokenizer: tokenizer
523+
tokenizer: tokenizer,
524+
additionalEosTokenIds: additionalEosTokenIds
509525
)
510526

511527
// Check if logits are requested
@@ -590,7 +606,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
590606
samplingConfiguration: SamplingConfiguration,
591607
maxTokens: Int,
592608
actualInputTokens: Int,
593-
modelVocabSize: Int?
609+
modelVocabSize: Int?,
610+
additionalEosTokenIds: [Int32] = []
594611
) async throws {
595612
let schema: String
596613
if FileManager.default.fileExists(atPath: schemaInput) {
@@ -603,7 +620,8 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
603620

604621
let stopSequences = try validateAndEncodeStopTokens(
605622
stopTokens: stopTokens,
606-
tokenizer: tokenizer
623+
tokenizer: tokenizer,
624+
additionalEosTokenIds: additionalEosTokenIds
607625
)
608626

609627
guard let vocabSize = modelVocabSize else {
@@ -676,15 +694,19 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
676694
/// - Parameters:
677695
/// - stopTokens: Array of stop token strings from CLI
678696
/// - tokenizer: Tokenizer to use for encoding
697+
/// - additionalEosTokenIds: Additional EOS token IDs from tokenizer config
679698
/// - Returns: StopSequences containing all valid sequences plus tokenizer EOS tokens
680699
func validateAndEncodeStopTokens(
681700
stopTokens: [String],
682-
tokenizer: any Tokenizer
701+
tokenizer: any Tokenizer,
702+
additionalEosTokenIds: [Int32] = []
683703
) throws -> StopSequences {
684704
var sequences: [[Int32]] = []
685705

686706
for stopString in stopTokens {
687-
let tokens = tokenizer.encode(text: stopString).map { Int32($0) }
707+
// Encode without adding BOS/EOS so special token strings like
708+
// "<end_of_turn>" resolve to their single token ID, not [BOS, id].
709+
let tokens = tokenizer.encode(text: stopString, addSpecialTokens: false).map { Int32($0) }
688710

689711
// Fatal error for empty encodings - user explicitly requested this stop token
690712
guard !tokens.isEmpty else {
@@ -710,7 +732,11 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
710732
}
711733

712734
// Use new initializer that automatically includes EOS tokens from tokenizer
713-
return StopSequences(for: tokenizer, additionalSequences: sequences)
735+
return StopSequences(
736+
for: tokenizer,
737+
additionalSequences: sequences,
738+
additionalEosTokenIds: additionalEosTokenIds
739+
)
714740
}
715741

716742
// MARK: - Asset Type Label

0 commit comments

Comments
 (0)