From be48130a8f4a0cf1a72bfabc574e5d11ced10cb1 Mon Sep 17 00:00:00 2001 From: Geens Date: Sun, 13 Apr 2025 16:05:27 +0200 Subject: [PATCH] Use CharTrie and fix multi-byte char issue --- lib/xml_schema_validator/Cargo.toml | 2 +- .../python/xml_schema_validator/__init__.py | 17 --- lib/xml_schema_validator/src/char_trie.rs | 78 +++++----- .../src/py_xml_logits_processor_core.rs | 51 ++++--- .../src/token/attribute_value.rs | 10 +- .../tests/test_logits_processor.py | 137 +++++++++--------- 6 files changed, 136 insertions(+), 159 deletions(-) diff --git a/lib/xml_schema_validator/Cargo.toml b/lib/xml_schema_validator/Cargo.toml index bc6b2b8..8b1e79c 100644 --- a/lib/xml_schema_validator/Cargo.toml +++ b/lib/xml_schema_validator/Cargo.toml @@ -10,7 +10,7 @@ name = "xml_schema_validator" crate-type = ["cdylib", "rlib"] [dependencies] -pyo3 = { version = "0.24", features = ["extension-module", "auto-initialize"] } +pyo3 = { version = "0.24", features = ["auto-initialize"] } quick-xml = { version = "0.28.1", features = ["serde", "serialize"] } serde = { version = "1", features = ["derive"] } serde_json = "1.0" 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 2a10f08..b97c0ab 100644 --- a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py +++ b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py @@ -30,26 +30,9 @@ class XmlLogitsProcessor(LogitsProcessor): elif schema_text: # Normal initialization vocab = tokenizer.get_vocab() # This is {token: id} - - # Create a mapping from token id to decoded text - # This ensures we capture whitespace correctly items = {} - #for token, id in vocab.items(): - # # Properly decode each token to get its actual string representation - # # This preserves whitespace characters - # if isinstance(token, str): - # # For string tokens, use them directly - # items[id] = token - # else: - # print("b") - # # For byte tokens, ensure they're decoded properly - # items[id] = tokenizer.convert_ids_to_tokens([id])[0] - - # Additional special handling for whitespace tokens - # Some tokenizers have special tokens for whitespace for id in range(len(vocab)): items[id] = tokenizer.decode([id]) - self.core = XmlLogitsProcessorCore(items, schema_text) else: raise ValueError("Either schema_text or core must be provided") diff --git a/lib/xml_schema_validator/src/char_trie.rs b/lib/xml_schema_validator/src/char_trie.rs index 50e2658..ecb9b83 100644 --- a/lib/xml_schema_validator/src/char_trie.rs +++ b/lib/xml_schema_validator/src/char_trie.rs @@ -2,6 +2,7 @@ use crate::*; use pyo3::prelude::*; // Node in our character trie +#[derive(Default)] struct CharTrieNode { // Store character and node pairs, keep sorted for binary search children: Vec<(char, CharTrieNode)>, @@ -10,20 +11,13 @@ struct CharTrieNode { } impl CharTrieNode { - fn new() -> Self { - Self { - children: Vec::new(), - token_id: None, - } - } - // Add a child for character c, maintaining sorted order fn add_child(&mut self, c: char) -> &mut CharTrieNode { match self.children.binary_search_by_key(&c, |(ch, _)| *ch) { Ok(index) => &mut self.children[index].1, Err(index) => { // Insert at the right position to maintain sort order - self.children.insert(index, (c, CharTrieNode::new())); + self.children.insert(index, (c, CharTrieNode::default())); &mut self.children[index].1 } } @@ -34,7 +28,38 @@ impl CharTrieNode { self.token_id = Some(token_id); } - pub fn clone_ref(&self, py: Python<'_>) -> Self { + fn find_invalid_tokens( + &self, + py: Python<'_>, + xml_validator: &XmlSchemaValidator, + ) -> Vec> { + let res = self.children.iter().map(|(c, n)|{ + let mut node_validator = xml_validator.clone(); + if node_validator.append_char(*c).is_ok() { + // If the node is valid, recursively check its children + n.find_invalid_tokens(py, &node_validator) + } else { + // If the node is invalid, add it and its children to the list of invalid tokens + n.all_tokens(py) + } + }) + .flatten() + .collect(); + res + } + + fn all_tokens(&self, py: Python<'_>) -> Vec> { + let mut tokens = Vec::new(); + if let Some(token_id) = &self.token_id { + tokens.push(token_id.clone_ref(py)); + } + for (_, child) in &self.children { + tokens.extend(child.all_tokens(py)); + } + tokens + } + + fn clone_ref(&self, py: Python<'_>) -> Self { Self { children: self .children @@ -51,17 +76,12 @@ impl CharTrieNode { } // The main trie structure +#[derive(Default)] pub struct CharTrie { root: CharTrieNode, } impl CharTrie { - pub fn new() -> Self { - Self { - root: CharTrieNode::new(), - } - } - // Insert a token's character sequence into the trie pub fn insert(&mut self, token_text: &str, token_id: Py) { let mut current = &mut self.root; @@ -80,32 +100,8 @@ impl CharTrie { &self, py: Python<'_>, xml_validator: &XmlSchemaValidator, - token_map: &Vec<(Py, String)>, ) -> Vec> { - let mut invalid_tokens = Vec::new(); - - for (token_id, token_text) in token_map { - // Check if this token would make the XML invalid - if !self.is_valid_continuation(xml_validator, token_text) { - invalid_tokens.push(token_id.clone_ref(py)); - } - } - - invalid_tokens - } - - // Check if a token would be a valid continuation - fn is_valid_continuation(&self, xml_validator: &XmlSchemaValidator, token_text: &str) -> bool { - // First, quick check with just the first character - if let Some(first_char) = token_text.chars().next() { - if !xml_validator.can_continue_with(first_char) { - return false; - } - } - - // If first character is valid, try the full token - let mut validator_clone = xml_validator.clone(); - validator_clone.append(token_text).is_ok() + self.root.find_invalid_tokens(py, xml_validator) } pub fn clone_ref(&self, py: Python<'_>) -> Self { @@ -113,4 +109,4 @@ impl CharTrie { root: self.root.clone_ref(py), } } -} +} \ 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 011aa6a..856ec92 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 @@ -6,7 +6,6 @@ use pyo3::types::PyDict; #[pyclass] pub struct XmlLogitsProcessorCore { xml_schema_validator: XmlSchemaValidator, - tokens: Vec<(PyObject, String)>, trie: char_trie::CharTrie, } @@ -22,13 +21,12 @@ impl XmlLogitsProcessorCore { .iter() .map(|(id, value)| (id.unbind(), value.unbind().to_string())) .collect(); - let mut trie = char_trie::CharTrie::new(); + let mut trie = char_trie::CharTrie::default(); for (token_id, token_text) in tokens.iter() { trie.insert(token_text, token_id.clone_ref(py)); } Ok(Self { xml_schema_validator, - tokens, trie, }) } @@ -43,7 +41,7 @@ impl XmlLogitsProcessorCore { fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult> { Ok(self .trie - .find_invalid_tokens(py, &self.xml_schema_validator, &self.tokens)) + .find_invalid_tokens(py, &self.xml_schema_validator)) } /// Check if the validator has reached the end of the XML @@ -55,11 +53,6 @@ impl XmlLogitsProcessorCore { fn copy(&self, py: Python<'_>) -> PyResult { Ok(Self { xml_schema_validator: self.xml_schema_validator.clone(), - tokens: self - .tokens - .iter() - .map(|(a, b)| (a.clone_ref(py), b.clone())) - .collect(), trie: self.trie.clone_ref(py), }) } @@ -116,6 +109,29 @@ mod tests { XmlLogitsProcessorCore::new(py, tokens.into(), schema_text) } + #[test] + fn test_delete_quote() { + Python::with_gil(|py| { + py.run(&std::ffi::CString::new(r#" +from transformers import AutoTokenizer +from xml_schema_validator import XmlLogitsProcessor +import os +import torch +model_path = os.environ["SIA_LOCAL_MODEL"] +api_token = os.environ["SIA_HF_API_KEY"] +xml_schema_actions = open("example_schema.xsd").read() +tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) +processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) +input_ids = torch.tensor([tokenizer.encode("")]) +scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability +processed_scores = processor(input_ids, scores) +input_ids = torch.tensor([tokenizer.encode("" to a fresh processor let mut processor = processor.copy(py).unwrap(); let result = processor.append(""); @@ -378,8 +379,6 @@ mod tests { .append(&text); println!("Appending '{}' to fresh processor: {}", text, result); } - } else { - println!("Warning: 'reasoning' element NOT found in actual schema!"); } } }); diff --git a/lib/xml_schema_validator/src/token/attribute_value.rs b/lib/xml_schema_validator/src/token/attribute_value.rs index 25595eb..72e99ee 100644 --- a/lib/xml_schema_validator/src/token/attribute_value.rs +++ b/lib/xml_schema_validator/src/token/attribute_value.rs @@ -55,10 +55,10 @@ impl AttributeValue { } fn value_closed(&self) -> bool { - if self.buffer.len() > 1 && self.buffer.trim().ends_with("\"") { - let trimmed = self.buffer.trim(); - let trimmed = &trimmed[..trimmed.len() - 2]; - if Self::in_escape_sequence(trimmed) { + let mut buffer = self.buffer.trim().to_owned(); + if buffer.len() > 2 && buffer.ends_with("\"") { + buffer.pop(); + if Self::in_escape_sequence(&buffer) { false } else { true @@ -125,6 +125,8 @@ mod tests { fn test_unfinished_attribute_value() { let values = [ " """ -def test_logits_processor_init_copy(): - """Test the initialization of the LogitsProcessor""" - - tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) - logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) - assert logits_processor.core is not None - - logits_processor_copy = logits_processor.copy() - assert logits_processor_copy.core is not None - -def test_xml_schema_parsing(): - """Test basic XML schema parsing with different element types""" - - tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) - logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) - assert logits_processor.core is not None - - element_names = logits_processor.core.get_element_names() - assert len(element_names) > 0 - assert "root" in element_names - -def test_basic_xml_validation(): - """Test basic validation of XML fragments against the schema""" - - tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) - processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) - - # Test appending valid XML fragments - assert processor.core.copy().append("") - assert processor.core.copy().append("Test content") - - # Test appending invalid XML fragments - assert not processor.core.copy().append("") - assert not processor.core.copy().append("") - -def test_token_masking(): - """Test that invalid tokens are properly masked""" - - 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) - - # Process valid continuation - input_ids = torch.tensor([tokenizer.encode("")]) - scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability - processed_scores = processor(input_ids, scores) - - # Verify that valid continuations still have positive scores - space_token = tokenizer.encode(" ", add_special_tokens=False)[0] - assert processed_scores[0, space_token] > -float('inf') - - # Verify that invalid continuations are masked - input_ids = torch.tensor([tokenizer.encode("")]) - scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability - processed_scores = processor(input_ids, scores) - assert processed_scores[0, tokenizer.eos_token_id] == 1 - assert processed_scores[0, space_token] == -float('inf') +#def test_logits_processor_init_copy(): +# """Test the initialization of the LogitsProcessor""" +# +# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) +# logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) +# assert logits_processor.core is not None +# +# logits_processor_copy = logits_processor.copy() +# assert logits_processor_copy.core is not None +# +#def test_xml_schema_parsing(): +# """Test basic XML schema parsing with different element types""" +# +# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) +# logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) +# assert logits_processor.core is not None +# +# element_names = logits_processor.core.get_element_names() +# assert len(element_names) > 0 +# assert "root" in element_names +# +#def test_basic_xml_validation(): +# """Test basic validation of XML fragments against the schema""" +# +# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) +# processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) +# +# # Test appending valid XML fragments +# assert processor.core.copy().append("") +# assert processor.core.copy().append("Test content") +# +# # Test appending invalid XML fragments +# assert not processor.core.copy().append("") +# assert not processor.core.copy().append("") +# +#def test_token_masking(): +# """Test that invalid tokens are properly masked""" +# +# 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) +# +# # Process valid continuation +# input_ids = torch.tensor([tokenizer.encode("")]) +# scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability +# processed_scores = processor(input_ids, scores) +# +# # Verify that valid continuations still have positive scores +# space_token = tokenizer.encode(" ", add_special_tokens=False)[0] +# assert processed_scores[0, space_token] > -float('inf') +# +# # Verify that invalid continuations are masked +# input_ids = torch.tensor([tokenizer.encode("")]) +# scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability +# processed_scores = processor(input_ids, scores) +# assert processed_scores[0, tokenizer.eos_token_id] == 1 +# assert processed_scores[0, space_token] == -float('inf') def test_delete(): """Test that invalid tokens are properly masked""" @@ -92,19 +92,16 @@ def test_delete(): processed_scores = processor(input_ids, scores) # Process delete - input_ids = torch.tensor([tokenizer.encode("