Add backwards-compatible support for multiple EOS tokens

c_sample/c_sample.cpp

@@ -29,6 +29,10 @@ LlgTokenizer *create_tokenizer(std::vector<std::vector<uint8_t>> &tokens,
   LlgTokenizerInit tok_init = {};
   tok_init.vocab_size = (uint32_t)tokens.size();
   tok_init.tok_eos = tok_eos;
+  // For models with multiple EOS tokens (e.g., Qwen3), set:
+  //   LlgToken extra_eos[] = {second_eos, third_eos};
+  //   tok_init.tok_eos_extra = extra_eos;
+  //   tok_init.tok_eos_extra_count = sizeof(extra_eos) / sizeof(extra_eos[0]);
   tok_init.token_lens = token_lens;
   tok_init.token_bytes = token_bytes;
   tok_init.tokenize_assumes_string = false;

parser/llguidance.h

@@ -239,6 +239,16 @@ typedef struct LlgTokenizerInit {
    * Pass NULL to use defaults. Pass empty array to disable.
    */
   const char *const *slices;
+  /**
+   * Additional EOS token IDs beyond `tok_eos`.
+   * Points to an array of `tok_eos_extra_count` elements.
+   * When NULL (the default for zero-initialized structs), only `tok_eos` is used.
+   */
+  const LlgToken *tok_eos_extra;
+  /**
+   * Number of elements in the `tok_eos_extra` array.
+   */
+  uint32_t tok_eos_extra_count;
 } LlgTokenizerInit;
 
 

parser/src/ffi.rs

@@ -137,7 +137,16 @@ impl LlgTokenizer {
             token_bytes
         };
 
-        let trie = TokTrie::from(&TokRxInfo::new(tokens.len() as u32, init.tok_eos), &tokens);
+        let mut trie = TokTrie::from(&TokRxInfo::new(tokens.len() as u32, init.tok_eos), &tokens);
+
+        if !init.tok_eos_extra.is_null() && init.tok_eos_extra_count > 0 {
+            let extra = unsafe {
+                std::slice::from_raw_parts(init.tok_eos_extra, init.tok_eos_extra_count as usize)
+            };
+            let mut eos_tokens = vec![init.tok_eos];
+            eos_tokens.extend_from_slice(extra);
+            trie = trie.with_eos_tokens(&eos_tokens);
+        }
 
         let tok_env: TokEnv = Arc::new(CTokenizerInner {
             trie,
@@ -249,6 +258,14 @@ pub struct LlgTokenizerInit {
     /// This is array of pointers to strings, terminated with NULL (argv style).
     /// Pass NULL to use defaults. Pass empty array to disable.
     pub slices: *const *const c_char,
+
+    /// Additional EOS token IDs beyond `tok_eos`.
+    /// Points to an array of `tok_eos_extra_count` elements.
+    /// When NULL (the default for zero-initialized structs), only `tok_eos` is used.
+    pub tok_eos_extra: *const LlgToken,
+
+    /// Number of elements in the `tok_eos_extra` array.
+    pub tok_eos_extra_count: u32,
 }
 
 #[derive(Clone)]

parser/src/tokenparser.rs

@@ -21,7 +21,7 @@ pub struct TokenParser {
     pub dbg_grammar: String,
     last_step_stats: ParserStats,
     max_step_stats: ParserStats,
-    eos_token: TokenId,
+    eos_tokens: Vec<TokenId>,
 
     had_rollback: bool,
     had_backtrack: bool,
@@ -91,7 +91,7 @@ impl TokenParser {
             factory.perf_counters(),
         )?;
         parser.metrics_mut().rand = factory.next_rng();
-        let eos_token = token_env.tok_trie().eos_token();
+        let eos_tokens = token_env.tok_trie().eos_tokens().to_vec();
 
         Ok(TokenParser {
             bias_computer: factory.slicer().clone(),
@@ -108,7 +108,7 @@ impl TokenParser {
             error_message: None,
             parser,
             dbg_grammar: String::new(),
-            eos_token,
+            eos_tokens,
             llm_tokens: Vec::new(),
             llm_bytes: Vec::new(),
             grm_prefix: Vec::new(),
@@ -389,7 +389,7 @@ impl TokenParser {
         let new_len = self.llm_tokens.len() - n_tokens;
         let mut bytes_to_drop = 0;
         for tok in &self.llm_tokens[new_len..] {
-            if *tok == self.eos_token {
+            if self.eos_tokens.contains(tok) {
                 // doesn't count; we hope it's last though...
                 bytes_to_drop += 0;
             } else {
@@ -492,8 +492,12 @@ impl TokenParser {
             return Err(self.stop_for_parser_error("", s));
         }
 
-        if self.eos_token != INVALID_TOKEN && self.is_accepting() {
-            allowed_tokens.allow_token(self.eos_token);
+        if self.is_accepting() {
+            for &eos in &self.eos_tokens {
+                if eos != INVALID_TOKEN {
+                    allowed_tokens.allow_token(eos);
+                }
+            }
         }
 
         self.log_final(&prefix, &allowed_tokens);
@@ -797,7 +801,7 @@ impl TokenParser {
         }
         self.max_tokens_total -= 1;
 
-        if token == self.eos_token {
+        if self.eos_tokens.contains(&token) {
             if self.parser.scan_eos() {
                 // it got scanned correctly, so we remove it
                 // this only happens for gen() terminated by EOS
@@ -838,7 +842,10 @@ impl TokenParser {
     /// This generally should be called after consume_token().
     pub fn check_stop(&mut self) -> Result<bool> {
         let empty_token_prefix = !self.has_ff_bytes();
-        let pending_eos = self.llm_tokens.last() == Some(&self.eos_token);
+        let pending_eos = self
+            .llm_tokens
+            .last()
+            .is_some_and(|t| self.eos_tokens.contains(t));
         let lexer_bytes = self.parser.has_pending_lexeme_bytes();
         let is_accepting = self.is_accepting();
         let can_advance = self.parser.can_advance();

python/llguidance/_lib.pyi

@@ -6,13 +6,14 @@ from ._tokenizer import TokenizerWrapper
 class LLTokenizer:
     vocab_size: int
     eos_token: TokenId
+    eos_tokens: List[TokenId]
     is_canonical: bool
 
     def __new__(
         cls,
         tokenizer: Union[str, TokenizerWrapper],
         n_vocab: Optional[int] = None,
-        eos_token: Optional[TokenId] = None,
+        eos_token: Optional[Union[TokenId, List[TokenId]]] = None,
         slices: Optional[List[str]] = None,
     ) -> "LLTokenizer":
         """
@@ -23,6 +24,7 @@ class LLTokenizer:
         Args:
             tokenizer: str or TokenizerWrapper - if str, it is the name or path to the HF tokenizers tokenizer; otherwise it is a TokenizerWrapper
             n_vocab: int - override the size of the vocabulary
+            eos_token: int or list of ints - override the EOS token(s)
             slices: List[str] - configuration for slicer optimization; pass [] to disable,
                 or None to use general_slices()
         """

python/llguidance/hf.py

@@ -1,5 +1,5 @@
 from copy import copy
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import transformers
 
@@ -9,7 +9,7 @@
 def from_tokenizer(
     hf_tokenizer: transformers.PreTrainedTokenizerFast,
     n_vocab: Optional[int] = None,
-    eos_token: Optional[int] = None,
+    eos_token: Optional[Union[int, List[int]]] = None,
     slices: Optional[List[str]] = None,
 ) -> LLTokenizer:
     """
@@ -21,7 +21,7 @@ def from_tokenizer(
     Args:
         hf_tokenizer: transformers.PreTrainedTokenizerFast - the tokenizer to wrap
         n_vocab: int - override the size of the vocabulary
-        eos_token: int - override the EOS token
+        eos_token: int or list of ints - override the EOS token(s)
         slices: List[str] - configuration for slicer optimization; pass [] to disable,
             or None to use the default configuration
     """

python/llguidance/llamacpp.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Union
 
 from ._lib import LLTokenizer
 
@@ -8,7 +8,7 @@
 def lltokenizer_from_vocab(
     vocab: llama_cpp.llama_vocab_p,
     n_vocab: Optional[int] = None,
-    eos_token: Optional[int] = None,
+    eos_token: Optional[Union[int, List[int]]] = None,
     slices: Optional[List[str]] = None,
 ) -> LLTokenizer:
     """
@@ -18,6 +18,7 @@ def lltokenizer_from_vocab(
     Args:
         vocab: llama_cpp.llama_vocab_p - the vocab object to use
         n_vocab: int - override the size of the vocabulary
+        eos_token: int or list of ints - override the EOS token(s)
         eos_token: int - override the EOS token
         slices: List[str] - configuration for slicer optimization; pass [] to disable,
             or None to use the default configuration

python/llguidance/tiktoken.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, TYPE_CHECKING
+from typing import List, Optional, Union, TYPE_CHECKING
 
 from ._lib import LLTokenizer
 
@@ -10,7 +10,7 @@ def lltokenizer_from_encoding(
     encoding: 'tiktoken.Encoding',
     *,
     n_vocab: Optional[int] = None,
-    eos_token: Optional[int] = None,
+    eos_token: Optional[Union[int, List[int]]] = None,
     slices: Optional[List[str]] = None,
 ) -> LLTokenizer:
     """
@@ -20,7 +20,7 @@ def lltokenizer_from_encoding(
     Args:
         encoding: tiktoken.Encoding - the encoding object to use
         n_vocab: int - override the size of the vocabulary
-        eos_token: int - override the EOS token
+        eos_token: int or list of ints - override the EOS token(s)
         slices: List[str] - configuration for slicer optimization; pass [] to disable,
             or None to use the default configuration
     """

python_ext/src/llamatokenizer.rs

@@ -119,13 +119,17 @@ pub fn tokenv_from_llamacpp(
     tokens: Vec<Vec<u8>>,
     vocab_ptr: usize,
     tokenize_fptr: usize,
-    eos_token: u32,
+    eos_tokens: &[u32],
 ) -> Result<TokEnv> {
+    ensure!(!eos_tokens.is_empty(), "eos_tokens must not be empty");
     ensure!(vocab_ptr != 0, "vocab_ptr must be non-null");
     ensure!(tokenize_fptr != 0, "tokenize_fptr must be non-null");
 
-    let info = TokRxInfo::new(tokens.len() as u32, eos_token);
-    let trie = TokTrie::from(&info, &tokens);
+    let info = TokRxInfo::new(tokens.len() as u32, eos_tokens[0]);
+    let mut trie = TokTrie::from(&info, &tokens);
+    if eos_tokens.len() > 1 {
+        trie = trie.with_eos_tokens(eos_tokens);
+    }
 
     let mut llama_tok = LlamaTokenizer {
         trie,

python_ext/src/py.rs

@@ -16,6 +16,23 @@ use toktrie_tiktoken::TikTokenBPE;
 
 use crate::llamatokenizer::tokenv_from_llamacpp;
 
+/// Extract eos_token from a Python value that may be int or list[int].
+/// Returns None if the value is None, or Some(vec) if it's an int or list of ints.
+fn extract_eos_tokens(obj: &Bound<'_, PyAny>) -> PyResult<Vec<u32>> {
+    if let Ok(single) = obj.extract::<u32>() {
+        Ok(vec![single])
+    } else if let Ok(list) = obj.extract::<Vec<u32>>() {
+        if list.is_empty() {
+            return Err(PyValueError::new_err("eos_token list must not be empty"));
+        }
+        Ok(list)
+    } else {
+        Err(PyValueError::new_err(
+            "eos_token must be an int or a non-empty list of ints",
+        ))
+    }
+}
+
 struct PyTokenizer {
     tok_trie: Arc<toktrie::TokTrie>,
     tokenizer_fun: Py<PyAny>,
@@ -36,9 +53,10 @@ impl LLTokenizer {
     fn py_new(
         tokenizer: Bound<'_, PyAny>,
         n_vocab: Option<usize>,
-        eos_token: Option<u32>,
+        eos_token: Option<Bound<'_, PyAny>>,
         slices: Option<Vec<String>>,
     ) -> PyResult<Self> {
+        let eos_tokens = eos_token.as_ref().map(extract_eos_tokens).transpose()?;
         let tok_env: TokEnv = if let Ok(tokenizer_str) = tokenizer.extract::<String>() {
             if tokenizer_str == "byte" {
                 ApproximateTokEnv::single_byte_env()
@@ -48,8 +66,8 @@ impl LLTokenizer {
                 } else {
                     ByteTokenizer::from_file(&tokenizer_str).map_err(val_error)?
                 };
-                if let Some(eos_token) = eos_token {
-                    tok.set_eos_token(eos_token);
+                if let Some(ref eos_tokens) = eos_tokens {
+                    tok.set_eos_tokens(eos_tokens);
                 }
                 tok.into_tok_env(n_vocab).map_err(val_error)?
             }
@@ -77,18 +95,22 @@ impl LLTokenizer {
         encoder: HashMap<Vec<u8>, u32>,
         special_tokens: HashMap<String, u32>,
         pattern: &str,
-        eos_token: u32,
+        eos_token: Bound<'_, PyAny>,
         n_vocab: Option<usize>,
         slices: Option<Vec<String>>,
     ) -> PyResult<Self> {
-        let bpe = TikTokenBPE::new(
+        let eos_tokens = extract_eos_tokens(&eos_token)?;
+        let mut bpe = TikTokenBPE::new(
             encoder.into_iter().collect(),
             special_tokens.into_iter().collect(),
             pattern,
             n_vocab,
-            eos_token,
+            eos_tokens[0],
         )
         .map_err(val_error)?;
+        if eos_tokens.len() > 1 {
+            bpe.set_eos_tokens(&eos_tokens);
+        }
         let tok_env = bpe.to_env();
 
         let factory = ParserFactory::new(
@@ -108,11 +130,12 @@ impl LLTokenizer {
         tokens: Vec<Vec<u8>>,
         vocab_ptr: usize,
         tokenize_fptr: usize,
-        eos_token: u32,
+        eos_token: Bound<'_, PyAny>,
         slices: Option<Vec<String>>,
     ) -> PyResult<Self> {
-        let tok_env =
-            tokenv_from_llamacpp(tokens, vocab_ptr, tokenize_fptr, eos_token).map_err(val_error)?;
+        let eos_tokens = extract_eos_tokens(&eos_token)?;
+        let tok_env = tokenv_from_llamacpp(tokens, vocab_ptr, tokenize_fptr, &eos_tokens)
+            .map_err(val_error)?;
 
         let factory = ParserFactory::new(
             &tok_env,
@@ -244,6 +267,11 @@ impl LLTokenizer {
     fn eos_token(&self) -> u32 {
         self.tok_trie().eos_token()
     }
+
+    #[getter]
+    fn eos_tokens(&self) -> Vec<u32> {
+        self.tok_trie().eos_tokens().to_vec()
+    }
 }
 
 impl LLTokenizer {