allow for using -1 for eos token (if missing); untested

guidance-ai · Jan 8, 2025 · b40e1c7 · b40e1c7
1 parent 5198689
commit b40e1c7
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 6 deletions.
diff --git a/parser/src/earley/parser.rs b/parser/src/earley/parser.rs
@@ -16,7 +16,7 @@ use derivre::{AlphabetInfo, RegexAst, StateID};
 use hashbrown::HashSet;
 use instant::Instant;
 use serde::{Deserialize, Serialize};
-use toktrie::{Recognizer, SimpleVob, TokEnv, TokTrie};
+use toktrie::{Recognizer, SimpleVob, TokEnv, TokTrie, INVALID_TOKEN};
 
 use crate::{
     api::{ParserLimits, StopReason},
@@ -663,8 +663,9 @@ impl ParserState {
             let _ = self.flush_lexer();
         }
 
-        if start.is_empty() && self.lexer_allows_eos() {
-            set.allow_token(computer.trie().eos_token());
+        let eos = computer.trie().eos_token();
+        if eos != INVALID_TOKEN && start.is_empty() && self.lexer_allows_eos() {
+            set.allow_token(eos);
         }
 
         self.stats.compute_time_us += t0.elapsed().as_micros() as u64;

diff --git a/parser/src/tokenparser.rs b/parser/src/tokenparser.rs
@@ -9,7 +9,7 @@ use crate::{
 };
 use anyhow::{ensure, Result};
 use serde_json::json;
-use toktrie::{InferenceCapabilities, SimpleVob, TokEnv, TokenId};
+use toktrie::{InferenceCapabilities, SimpleVob, TokEnv, TokenId, INVALID_TOKEN};
 
 #[derive(Clone)]
 pub struct TokenParser {
@@ -386,7 +386,7 @@ impl TokenParser {
             return Err(self.stop_for_parser_error("", s));
         }
 
-        if self.is_accepting() {
+        if self.eos_token != INVALID_TOKEN && self.is_accepting() {
             allowed_tokens.allow_token(self.eos_token);
         }
 

diff --git a/toktrie/src/lib.rs b/toktrie/src/lib.rs
@@ -9,6 +9,7 @@ mod toktree;
 pub use svob::{SimpleVob, SimpleVobIter};
 pub use toktree::{
     Recognizer, TokEnv, TokEnvWithTrie, TokRxInfo, TokTrie, TokenId, TokenizerEnv, TrieNode,
+    INVALID_TOKEN,
 };
 
 /// Defines what is allowed in Branch

diff --git a/toktrie/src/toktree.rs b/toktrie/src/toktree.rs
@@ -211,6 +211,8 @@ pub struct TrieNode {
     bits2: u32,
 }
 
+pub const INVALID_TOKEN: TokenId = 0xffff_ffff;
+
 const NO_TOKEN: u32 = 0xffffff;
 
 impl TrieNode {
@@ -345,7 +347,7 @@ impl TokTrie {
         let max_tok = std::cmp::min(max_examples, num_set);
         let mut token_names = Vec::new();
         // make sure we include EOS first if it's allowed
-        if ts1.is_allowed(self.info.tok_eos) {
+        if self.info.tok_eos != INVALID_TOKEN && ts1.is_allowed(self.info.tok_eos) {
             token_names.push("EOS".to_string());
         }
         for idx in 0..self.vocab_size() {