Fix char_trie to store all token ids

lib/xml_schema_validator/python/xml_schema_validator/__init__.py

@@ -1,5 +1,4 @@
 import torch
-import sys
 from transformers import LogitsProcessor
 from transformers import AutoTokenizer
 from xml_schema_validator._rs import XmlLogitsProcessorCore
@@ -78,20 +77,6 @@ class XmlLogitsProcessor(LogitsProcessor):
             
             # Find the last occurrence of the assistant start marker
             last_marker_pos = full_text.rfind(self.assistant_start_marker)
-            
-            if last_marker_pos == -1:
-                # If primary marker not found, try looking for a second common marker as fallback
-                fallback_markers = ["<|assistant|>", "\nAssistant: ", "\nA: "]
-                for fallback in fallback_markers:
-                    last_marker_pos = full_text.rfind(fallback)
-                    if last_marker_pos != -1:
-                        # Found a fallback marker
-                        self.assistant_start_marker = fallback  # Update for future calls
-                        break
-                        
-                # If still not found, we can't determine where the assistant content starts
-                if last_marker_pos == -1:
-                    continue
                 
             # Extract the assistant's response by taking text after the marker plus its length
             start_pos = last_marker_pos + len(self.assistant_start_marker)

lib/xml_schema_validator/src/char_trie.rs

@@ -7,8 +7,8 @@ use rayon::prelude::*;
 struct CharTrieNode {
     // Store character and node pairs, keep sorted for binary search
     children: Vec<(char, CharTrieNode)>,
-    // Store token IDs that are valid at this point
+    // Store token IDs that are valid at this point - CHANGED to Vec<u32>
-    token_id: Option<Arc<Py<PyAny>>>,
+    token_ids: Vec<u32>,
 }
 
 impl CharTrieNode {
@@ -24,15 +24,18 @@ impl CharTrieNode {
         }
     }
 
-    // Add a token id to this node
+    // Add a token id to this node - CHANGED to handle Vec<u32>
-    fn add_token(&mut self, token_id: Arc<Py<PyAny>>) {
+    fn add_token(&mut self, token_id: u32) {
-        self.token_id = Some(token_id);
+        // Only add if not already present (avoid duplicates)
+        if !self.token_ids.contains(&token_id) {
+            self.token_ids.push(token_id);
+        }
     }
 
     fn find_invalid_tokens(
         &self,
         xml_validator: &XmlSchemaValidator,
-    ) -> Vec<Arc<Py<PyAny>>> {
+    ) -> Vec<u32> {
         let res = self
             .children
             .iter()
@@ -55,7 +58,7 @@ impl CharTrieNode {
     fn find_invalid_tokens_parallel(
         &self,
         xml_validator: &XmlSchemaValidator,
-    ) -> Vec<Arc<Py<PyAny>>> {
+    ) -> Vec<u32> {
         // Convert children to a parallel iterator
         self.children
             .par_iter() // Use parallel iterator from rayon
@@ -76,11 +79,13 @@ impl CharTrieNode {
             .collect()
     }
 
-    fn all_tokens(&self) -> Vec<Arc<Py<PyAny>>> {
+    // CHANGED to handle Vec<u32>
+    fn all_tokens(&self) -> Vec<u32> {
         let mut tokens = Vec::new();
-        if let Some(token_id) = &self.token_id {
+        // Add all token IDs stored at this node
-            tokens.push(token_id.clone());
+        tokens.extend_from_slice(&self.token_ids);
-        }
+        
+        // Recursively collect token IDs from all children
         for (_, child) in &self.children {
             tokens.extend(child.all_tokens());
         }
@@ -96,7 +101,7 @@ pub struct CharTrie {
 
 impl CharTrie {
     // Insert a token's character sequence into the trie
-    pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) {
+    pub fn insert(&mut self, token_text: &str, token_id: u32) {
         let mut current = &mut self.root;
 
         // Add each character in the token text
@@ -105,15 +110,20 @@ impl CharTrie {
         }
 
         // Mark this node as representing token_id
-        current.add_token(Arc::new(token_id));
+        current.add_token(token_id);
     }
 
     // Find all tokens that could be invalid given the current XML state
     pub fn find_invalid_tokens(
         &self,
         xml_validator: &XmlSchemaValidator,
-    ) -> Vec<Arc<Py<PyAny>>> {
+    ) -> Vec<u32> {
         // Use the parallel version starting at the root
         self.root.find_invalid_tokens_parallel(xml_validator)
     }
+    
+    // NEW: Return all token IDs registered in the trie
+    pub fn all_registered_tokens(&self) -> Vec<u32> {
+        self.root.all_tokens()
+    }
 }

lib/xml_schema_validator/src/py_xml_logits_processor_core.rs

@@ -8,6 +8,7 @@ use pyo3::types::PyDict;
 pub struct XmlLogitsProcessorCore {
     xml_schema_validator: XmlSchemaValidator,
     trie: char_trie::CharTrie,
+    tokens: Vec<(u32, String)>,
 }
 
 #[pymethods]
@@ -18,17 +19,21 @@ impl XmlLogitsProcessorCore {
             Ok(xml_schema_validator) => xml_schema_validator,
             Err(e) => return Err(pyo3::exceptions::PyValueError::new_err(e.to_string())),
         };
-        let tokens: Vec<(Py<PyAny>, String)> = tokens
+        let tokens: Vec<(u32, String)> = tokens
             .iter()
-            .map(|(id, value)| (id.unbind(), value.unbind().to_string()))
+            .map(|(id, value)| {
+                let id: u32 = id.extract().unwrap();
+                (id, value.unbind().to_string())
+            })
             .collect();
         let mut trie = char_trie::CharTrie::default();
         for (token_id, token_text) in tokens.iter() {
-            trie.insert(token_text, token_id.clone_ref(py));
+            trie.insert(token_text, token_id.clone());
         }
         Ok(Self {
             xml_schema_validator,
             trie,
+            tokens,
         })
     }
 
@@ -39,13 +44,10 @@ impl XmlLogitsProcessorCore {
     }
 
     /// Get a list of tokens that would lead to invalid XML
-    fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
+    fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<u32>> {
         Ok(self
             .trie
-            .find_invalid_tokens(&self.xml_schema_validator)
+            .find_invalid_tokens(&self.xml_schema_validator))
-            .iter()
-            .map(|token_id| token_id.clone_ref(py))
-            .collect())
     }
 
     /// Check if the validator has reached the end of the XML

lib/xml_schema_validator/src/token/mod.rs

@@ -49,6 +49,9 @@ pub enum Token {
 
 impl Token {
     pub fn append(self, c: char) -> Vec<Token> {
+        if !c.is_ascii_graphic() && ! c.is_ascii_whitespace() {
+            return vec![];
+        }
         match self {
             Token::AttributeEquals(attribute_equals) => attribute_equals.append(c),
             Token::AttributeName(attribute_name) => attribute_name.append(c),

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -3,7 +3,7 @@ from xml_schema_validator import XmlLogitsProcessor
 import os
 import torch
 
-model_path = "../../model"
+model_path = "/root/models/current"
 api_token = os.environ["SIA_HF_API_KEY"]
 
 xml_schema_actions = open("example_schema.xsd").read()
@@ -58,12 +58,32 @@ def test_token_masking():
     
     tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
     processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
-
+    
     # Process empty input to set prompt length
     input_ids = torch.tensor([tokenizer.encode("")])
     scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
     processed_scores = processor(input_ids, scores)
     
+    # Check if '<' is the only allowed token
+    less_than_token = tokenizer.encode("<", add_special_tokens=False)[0]
+    
+    # Count how many tokens have scores not equal to -inf
+    valid_tokens = (processed_scores[0] > -float('inf')).sum().item()
+    
+    # Check if the '<' token is allowed (not masked)
+    is_less_than_allowed = processed_scores[0, less_than_token] > -float('inf')
+    
+    print(f"Number of allowed tokens: {valid_tokens}")
+    print(f"Is '<' token allowed: {is_less_than_allowed}")
+    
+    # Always print the first 10 allowed tokens
+    allowed_token_ids = torch.where(processed_scores[0] > -float('inf'))[0].tolist()
+    # Take only the first 10 tokens (or all if less than 10)
+    allowed_token_ids_sample = allowed_token_ids[:min(10, len(allowed_token_ids))]
+    allowed_tokens = [tokenizer.decode([token_id]) for token_id in allowed_token_ids_sample]
+    print(f"First 10 allowed token IDs: {allowed_token_ids_sample}")
+    print(f"First 10 allowed tokens: {allowed_tokens}")
+    
     # Process valid continuation
     input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
     scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability

sia/web/response_websocket.py

@@ -1,4 +1,4 @@
-from aiohttp import web
+from aiohttp import web, WSMsgType
 
 from ..web_agent import WebAgent
 from .util import wrap_async