Fix eos token logic

Author: mattt

Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift

@@ -974,6 +974,7 @@ import Foundation
                     batch: batchPointer,
                     position: initialPosition,
                     maximumTokens: maxTokens,
+                    endTokens: [],
                     tokenToTextFn: { [self] token in self.tokenToText(vocab: vocab, token: llama_token(token)) }
                 )
                 var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema)
@@ -1004,6 +1005,7 @@ import Foundation
                 batch: UnsafeMutablePointer<llama_batch>,
                 position: Int32,
                 maximumTokens: Int,
+                endTokens: Set<Int>? = nil,
                 tokenToTextFn: @escaping (Int) -> String?
             ) {
                 self.context = context
@@ -1016,9 +1018,13 @@ import Foundation
                 self.totalTokenBudget = maximumTokens
                 self.eosToken = Int(llama_vocab_eos(vocab))
 
-                let eotTokenValue = llama_vocab_eot(vocab)
-                let endOfTurnToken = eotTokenValue != LLAMA_TOKEN_NULL ? Int(eotTokenValue) : eosToken
-                self.endTokens = [self.eosToken, endOfTurnToken]
+                if let endTokens {
+                    self.endTokens = endTokens
+                } else {
+                    let eotTokenValue = llama_vocab_eot(vocab)
+                    let endOfTurnToken = eotTokenValue != LLAMA_TOKEN_NULL ? Int(eotTokenValue) : eosToken
+                    self.endTokens = [self.eosToken, endOfTurnToken]
+                }
 
                 self.tokenToTextFn = tokenToTextFn
                 self.tokensExcludedFromRepetitionPenalty = Self.buildTokensExcludedFromRepetitionPenalty(

Sources/AnyLanguageModel/Models/MLXLanguageModel.swift

@@ -726,7 +726,8 @@ import Foundation
             context: context,
             input: lmInput,
             parameters: generateParameters,
-            maximumTokens: maxTokens
+            maximumTokens: maxTokens,
+            endTokens: []
         )
 
         var generator = try ConstrainedJSONGenerator(backend: backend, schema: schema)
@@ -805,7 +806,8 @@ import Foundation
             context: ModelContext,
             input: MLXLMCommon.LMInput,
             parameters: MLXLMCommon.GenerateParameters,
-            maximumTokens: Int
+            maximumTokens: Int,
+            endTokens: Set<Int>? = nil
         ) throws {
             self.model = context.model
             self.tokenizer = context.tokenizer
@@ -819,11 +821,15 @@ import Foundation
                 throw StructuredGenerationError.invalidVocabSize
             }
             self.eosToken = eosTokenId
-            self.endTokens = Self.buildEndTokens(
-                eosTokenId: eosTokenId,
-                tokenizer: context.tokenizer,
-                configuration: context.configuration
-            )
+            if let endTokens {
+                self.endTokens = endTokens
+            } else {
+                self.endTokens = Self.buildEndTokens(
+                    eosTokenId: eosTokenId,
+                    tokenizer: context.tokenizer,
+                    configuration: context.configuration
+                )
+            }
 
             self.tokensExcludedFromRepetitionPenalty = Self.buildTokensExcludedFromRepetitionPenalty(
                 tokenizer: context.tokenizer

Sources/AnyLanguageModel/StructuredGeneration.swift

@@ -51,7 +51,7 @@ struct ConstrainedJSONGenerator<Backend: TokenBackend> {
         let quoteToken = try Self.singleToken(for: "\"", backend: backend)
         self.quoteToken = quoteToken
 
-        self.stringTerminators = [quoteToken]
+        self.stringTerminators = backend.endTokens.union([quoteToken])
 
         var structuralTerminators = Set<Int>()
         for structuralText in [",", "}", "]", ":"] {