From 23b8c3c54d4415e2ee3af1ac583026633023f0f8 Mon Sep 17 00:00:00 2001 From: Geens Date: Fri, 18 Apr 2025 18:10:05 +0200 Subject: [PATCH] Fix CharTie multi-threading --- lib/xml_schema_validator/src/char_trie.rs | 95 +++++-------- .../src/py_xml_logits_processor_core.rs | 21 +-- .../tests/test_logits_processor.py | 130 +++++++++--------- 3 files changed, 108 insertions(+), 138 deletions(-) diff --git a/lib/xml_schema_validator/src/char_trie.rs b/lib/xml_schema_validator/src/char_trie.rs index 0674fbc..8d163cf 100644 --- a/lib/xml_schema_validator/src/char_trie.rs +++ b/lib/xml_schema_validator/src/char_trie.rs @@ -3,12 +3,12 @@ use pyo3::prelude::*; use rayon::prelude::*; // Node in our character trie -#[derive(Default)] +#[derive(Clone, Default)] 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>, + token_id: Option>>, } impl CharTrieNode { @@ -25,15 +25,14 @@ impl CharTrieNode { } // Add a token id to this node - fn add_token(&mut self, token_id: Py) { + fn add_token(&mut self, token_id: Arc>) { self.token_id = Some(token_id); } fn find_invalid_tokens( &self, - py: Python<'_>, xml_validator: &XmlSchemaValidator, - ) -> Vec> { + ) -> Vec>> { let res = self .children .iter() @@ -41,10 +40,10 @@ impl CharTrieNode { 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) + n.find_invalid_tokens(&node_validator) } else { // If the node is invalid, add it and its children to the list of invalid tokens - n.all_tokens(py) + n.all_tokens() } }) .flatten() @@ -55,65 +54,42 @@ impl CharTrieNode { // Parallel version of find_invalid_tokens for the first level of recursion fn find_invalid_tokens_parallel( &self, - _py: Python<'_>, xml_validator: &XmlSchemaValidator, - ) -> Vec> { - // We need to use the Python GIL to share PyAny objects across threads - Python::with_gil(|_py| { - // Convert children to a parallel iterator - self.children - .par_iter() // Use parallel iterator from rayon - .map(|(c, n)| { - // Clone validator for this thread - let mut node_validator = xml_validator.clone(); - - // Process this character - if node_validator.append_char(*c).is_ok() { - // For valid nodes, use the regular (non-parallel) method for deeper levels - Python::with_gil(|py| { - n.find_invalid_tokens(py, &node_validator) - }) - } else { - // For invalid nodes, get all tokens - Python::with_gil(|py| { - n.all_tokens(py) - }) - } - }) - .flatten() - .collect() - }) + ) -> Vec>> { + // Convert children to a parallel iterator + self.children + .par_iter() // Use parallel iterator from rayon + .map(|(c, n)| { + // Clone validator for this thread + let mut node_validator = xml_validator.clone(); + + // Process this character + if node_validator.append_char(*c).is_ok() { + // For valid nodes, use the regular (non-parallel) method for deeper levels + n.find_invalid_tokens(&node_validator) + } else { + // For invalid nodes, get all tokens + n.all_tokens() + } + }) + .flatten() + .collect() } - fn all_tokens(&self, py: Python<'_>) -> Vec> { + fn all_tokens(&self) -> Vec>> { let mut tokens = Vec::new(); if let Some(token_id) = &self.token_id { - tokens.push(token_id.clone_ref(py)); + tokens.push(token_id.clone()); } for (_, child) in &self.children { - tokens.extend(child.all_tokens(py)); + tokens.extend(child.all_tokens()); } tokens } - - fn clone_ref(&self, py: Python<'_>) -> Self { - Self { - children: self - .children - .iter() - .map(|(ch, node)| (*ch, node.clone_ref(py))) - .collect(), - token_id: self - .token_id - .iter() - .map(|token_id| token_id.clone_ref(py)) - .nth(0), - } - } } // The main trie structure -#[derive(Default)] +#[derive(Clone, Default)] pub struct CharTrie { root: CharTrieNode, } @@ -129,22 +105,15 @@ impl CharTrie { } // Mark this node as representing token_id - current.add_token(token_id); + current.add_token(Arc::new(token_id)); } // Find all tokens that could be invalid given the current XML state pub fn find_invalid_tokens( &self, - py: Python<'_>, xml_validator: &XmlSchemaValidator, - ) -> Vec> { + ) -> Vec>> { // Use the parallel version starting at the root - self.root.find_invalid_tokens_parallel(py, xml_validator) - } - - pub fn clone_ref(&self, py: Python<'_>) -> Self { - Self { - root: self.root.clone_ref(py), - } + self.root.find_invalid_tokens_parallel(xml_validator) } } \ 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 6217aed..d180f64 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 @@ -4,6 +4,7 @@ use pyo3::types::PyDict; /// A minimal LogitsProcessor that enforces valid XML according to a schema. #[pyclass] +#[derive(Clone)] pub struct XmlLogitsProcessorCore { xml_schema_validator: XmlSchemaValidator, trie: char_trie::CharTrie, @@ -41,7 +42,10 @@ impl XmlLogitsProcessorCore { fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult> { Ok(self .trie - .find_invalid_tokens(py, &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 @@ -50,11 +54,8 @@ impl XmlLogitsProcessorCore { } /// Create a copy of the processor - fn copy(&self, py: Python<'_>) -> PyResult { - Ok(Self { - xml_schema_validator: self.xml_schema_validator.clone(), - trie: self.trie.clone_ref(py), - }) + fn copy(&self) -> PyResult { + Ok(self.clone()) } /// Get a list of all element names in the schema @@ -258,7 +259,7 @@ processed_scores = processor(input_ids, scores) let processor = create_test_processor(py).unwrap(); // Make a copy - let copy = processor.copy(py).unwrap(); + let copy = processor.copy().unwrap(); // Original and copy should behave the same assert_eq!( @@ -268,7 +269,7 @@ processed_scores = processor(input_ids, scores) ); // Modify copy, original should be unchanged - let mut copy = processor.copy(py).unwrap(); + let mut copy = processor.copy().unwrap(); assert!(copy.append(""), "Copy should accept XML"); assert!(!processor.eof().unwrap(), "Original should be unchanged"); assert!(!copy.eof().unwrap(), "Copy should be modified"); @@ -345,7 +346,7 @@ processed_scores = processor(input_ids, scores) // Now try to append each character of "" one by one let chars = vec!['<', 'r', 'e', 'a', 's', 'o', 'n', 'i', 'n', 'g', '>']; let mut full_text = String::new(); - let processor = processor.copy(py).unwrap(); + let processor = processor.copy().unwrap(); for c in chars { full_text.push(c); @@ -373,7 +374,7 @@ processed_scores = processor(input_ids, scores) if element_names.contains(&"reasoning".to_string()) { // Test appending "" to a fresh processor - let mut processor = processor.copy(py).unwrap(); + let mut processor = processor.copy().unwrap(); let result = processor.append(""); println!("Appending '' to fresh processor: {}", result); diff --git a/lib/xml_schema_validator/tests/test_logits_processor.py b/lib/xml_schema_validator/tests/test_logits_processor.py index d20353d..67d59f9 100644 --- a/lib/xml_schema_validator/tests/test_logits_processor.py +++ b/lib/xml_schema_validator/tests/test_logits_processor.py @@ -18,67 +18,67 @@ xml_schema_only_root_node = """ """ -#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""" @@ -100,8 +100,8 @@ def test_delete(): if __name__ == "__main__": # Run individual tests for debugging - #test_logits_processor_init_copy() - #test_xml_schema_parsing() - #test_basic_xml_validation() - #test_token_masking() + test_logits_processor_init_copy() + test_xml_schema_parsing() + test_basic_xml_validation() + test_token_masking() test_delete()