From 5ac701316e9119635adeb650424703caf4656119 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Sat, 19 Apr 2025 11:49:09 +0000 Subject: [PATCH] Fix char_trie to store all token ids --- .../python/xml_schema_validator/__init__.py | 15 -------- lib/xml_schema_validator/src/char_trie.rs | 38 ++++++++++++------- .../src/py_xml_logits_processor_core.rs | 18 +++++---- lib/xml_schema_validator/src/token/mod.rs | 3 ++ .../tests/test_logits_processor.py | 24 +++++++++++- sia/web/response_websocket.py | 2 +- 6 files changed, 60 insertions(+), 40 deletions(-) diff --git a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py index d7651e7..22d78ac 100644 --- a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py +++ b/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) diff --git a/lib/xml_schema_validator/src/char_trie.rs b/lib/xml_schema_validator/src/char_trie.rs index 8d163cf..a0584fc 100644 --- a/lib/xml_schema_validator/src/char_trie.rs +++ b/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 - token_id: Option>>, + // Store token IDs that are valid at this point - CHANGED to Vec + token_ids: Vec, } impl CharTrieNode { @@ -24,15 +24,18 @@ impl CharTrieNode { } } - // Add a token id to this node - fn add_token(&mut self, token_id: Arc>) { - self.token_id = Some(token_id); + // Add a token id to this node - CHANGED to handle Vec + fn add_token(&mut self, token_id: u32) { + // 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>> { + ) -> Vec { let res = self .children .iter() @@ -55,7 +58,7 @@ impl CharTrieNode { fn find_invalid_tokens_parallel( &self, xml_validator: &XmlSchemaValidator, - ) -> Vec>> { + ) -> Vec { // 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>> { + // CHANGED to handle Vec + fn all_tokens(&self) -> Vec { let mut tokens = Vec::new(); - if let Some(token_id) = &self.token_id { - tokens.push(token_id.clone()); - } + // Add all token IDs stored at this node + 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) { + 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>> { + ) -> Vec { // 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 { + self.root.all_tokens() + } } \ No newline at end of file diff --git a/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs b/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs index d180f64..f6b0139 100644 --- a/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs +++ b/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, 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> { + fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult> { Ok(self .trie - .find_invalid_tokens(&self.xml_schema_validator) - .iter() - .map(|token_id| token_id.clone_ref(py)) - .collect()) + .find_invalid_tokens(&self.xml_schema_validator)) } /// Check if the validator has reached the end of the XML diff --git a/lib/xml_schema_validator/src/token/mod.rs b/lib/xml_schema_validator/src/token/mod.rs index 1e57efc..3921b20 100644 --- a/lib/xml_schema_validator/src/token/mod.rs +++ b/lib/xml_schema_validator/src/token/mod.rs @@ -49,6 +49,9 @@ pub enum Token { impl Token { pub fn append(self, c: char) -> Vec { + 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), diff --git a/lib/xml_schema_validator/tests/test_logits_processor.py b/lib/xml_schema_validator/tests/test_logits_processor.py index 67d59f9..a8c8621 100644 --- a/lib/xml_schema_validator/tests/test_logits_processor.py +++ b/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("")]) scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability diff --git a/sia/web/response_websocket.py b/sia/web/response_websocket.py index feb5b98..e72cce8 100644 --- a/sia/web/response_websocket.py +++ b/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