llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix CharTie multi-threading

Browse Source
This commit is contained in:
Geens
2025-04-18 18:10:05 +02:00
parent 64fc621d44
commit 23b8c3c54d
3 changed files with 108 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
lib/xml_schema_validator/src/char_trie.rs
View File

@@ -3,12 +3,12 @@ use pyo3::prelude::*;
use rayon::prelude::*; use rayon::prelude::*;
// Node in our character trie // Node in our character trie
#[derive(Default)] #[derive(Clone, Default)]
struct CharTrieNode { struct CharTrieNode {
// Store character and node pairs, keep sorted for binary search // Store character and node pairs, keep sorted for binary search
children: Vec<(char, CharTrieNode)>, children: Vec<(char, CharTrieNode)>,
// Store token IDs that are valid at this point // Store token IDs that are valid at this point
token_id: Option<Py<PyAny>>, token_id: Option<Arc<Py<PyAny>>>,
} }
impl CharTrieNode { impl CharTrieNode {
@@ -25,15 +25,14 @@ impl CharTrieNode {
} }
// Add a token id to this node // Add a token id to this node
fn add_token(&mut self, token_id: Py<PyAny>) { fn add_token(&mut self, token_id: Arc<Py<PyAny>>) {
self.token_id = Some(token_id); self.token_id = Some(token_id);
} }
fn find_invalid_tokens( fn find_invalid_tokens(
&self, &self,
py: Python<'_>,
xml_validator: &XmlSchemaValidator, xml_validator: &XmlSchemaValidator,
) -> Vec<Py<PyAny>> { ) -> Vec<Arc<Py<PyAny>>> {
let res = self let res = self
.children .children
.iter() .iter()
@@ -41,10 +40,10 @@ impl CharTrieNode {
let mut node_validator = xml_validator.clone(); let mut node_validator = xml_validator.clone();
if node_validator.append_char(*c).is_ok() { if node_validator.append_char(*c).is_ok() {
// If the node is valid, recursively check its children // If the node is valid, recursively check its children
n.find_invalid_tokens(py, &node_validator) n.find_invalid_tokens(&node_validator)
} else { } else {
// If the node is invalid, add it and its children to the list of invalid tokens // If the node is invalid, add it and its children to the list of invalid tokens
n.all_tokens(py) n.all_tokens()
} }
}) })
.flatten() .flatten()
@@ -55,65 +54,42 @@ impl CharTrieNode {
// Parallel version of find_invalid_tokens for the first level of recursion // Parallel version of find_invalid_tokens for the first level of recursion
fn find_invalid_tokens_parallel( fn find_invalid_tokens_parallel(
&self, &self,
_py: Python<'_>,
xml_validator: &XmlSchemaValidator, xml_validator: &XmlSchemaValidator,
) -> Vec<Py<PyAny>> { ) -> Vec<Arc<Py<PyAny>>> {
// We need to use the Python GIL to share PyAny objects across threads // Convert children to a parallel iterator
Python::with_gil(|_py| { self.children
// Convert children to a parallel iterator .par_iter() // Use parallel iterator from rayon
self.children .map(|(c, n)| {
.par_iter() // Use parallel iterator from rayon // Clone validator for this thread
.map(|(c, n)| { let mut node_validator = xml_validator.clone();
// Clone validator for this thread
let mut node_validator = xml_validator.clone(); // Process this character
if node_validator.append_char(*c).is_ok() {
// Process this character // For valid nodes, use the regular (non-parallel) method for deeper levels
if node_validator.append_char(*c).is_ok() { n.find_invalid_tokens(&node_validator)
// For valid nodes, use the regular (non-parallel) method for deeper levels } else {
Python::with_gil(|py| { // For invalid nodes, get all tokens
n.find_invalid_tokens(py, &node_validator) n.all_tokens()
}) }
} else { })
// For invalid nodes, get all tokens .flatten()
Python::with_gil(|py| { .collect()
n.all_tokens(py)
})
}
})
.flatten()
.collect()
})
} }
fn all_tokens(&self, py: Python<'_>) -> Vec<Py<PyAny>> { fn all_tokens(&self) -> Vec<Arc<Py<PyAny>>> {
let mut tokens = Vec::new(); let mut tokens = Vec::new();
if let Some(token_id) = &self.token_id { 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 { for (_, child) in &self.children {
tokens.extend(child.all_tokens(py)); tokens.extend(child.all_tokens());
} }
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 // The main trie structure
#[derive(Default)] #[derive(Clone, Default)]
pub struct CharTrie { pub struct CharTrie {
root: CharTrieNode, root: CharTrieNode,
} }
@@ -129,22 +105,15 @@ impl CharTrie {
} }
// Mark this node as representing token_id // 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 // Find all tokens that could be invalid given the current XML state
pub fn find_invalid_tokens( pub fn find_invalid_tokens(
&self, &self,
py: Python<'_>,
xml_validator: &XmlSchemaValidator, xml_validator: &XmlSchemaValidator,
) -> Vec<Py<PyAny>> { ) -> Vec<Arc<Py<PyAny>>> {
// Use the parallel version starting at the root // Use the parallel version starting at the root
self.root.find_invalid_tokens_parallel(py, xml_validator) self.root.find_invalid_tokens_parallel(xml_validator)
}
pub fn clone_ref(&self, py: Python<'_>) -> Self {
Self {
root: self.root.clone_ref(py),
}
} }
} }

21
lib/xml_schema_validator/src/py_xml_logits_processor_core.rs
View File

@@ -4,6 +4,7 @@ use pyo3::types::PyDict;
/// A minimal LogitsProcessor that enforces valid XML according to a schema. /// A minimal LogitsProcessor that enforces valid XML according to a schema.
#[pyclass] #[pyclass]
#[derive(Clone)]
pub struct XmlLogitsProcessorCore { pub struct XmlLogitsProcessorCore {
xml_schema_validator: XmlSchemaValidator, xml_schema_validator: XmlSchemaValidator,
trie: char_trie::CharTrie, trie: char_trie::CharTrie,
@@ -41,7 +42,10 @@ impl XmlLogitsProcessorCore {
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> { fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
Ok(self Ok(self
.trie .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 /// Check if the validator has reached the end of the XML
@@ -50,11 +54,8 @@ impl XmlLogitsProcessorCore {
} }
/// Create a copy of the processor /// Create a copy of the processor
fn copy(&self, py: Python<'_>) -> PyResult<Self> { fn copy(&self) -> PyResult<Self> {
Ok(Self { Ok(self.clone())
xml_schema_validator: self.xml_schema_validator.clone(),
trie: self.trie.clone_ref(py),
})
} }
/// Get a list of all element names in the schema /// 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(); let processor = create_test_processor(py).unwrap();
// Make a copy // Make a copy
let copy = processor.copy(py).unwrap(); let copy = processor.copy().unwrap();
// Original and copy should behave the same // Original and copy should behave the same
assert_eq!( assert_eq!(
@@ -268,7 +269,7 @@ processed_scores = processor(input_ids, scores)
); );
// Modify copy, original should be unchanged // Modify copy, original should be unchanged
let mut copy = processor.copy(py).unwrap(); let mut copy = processor.copy().unwrap();
assert!(copy.append("<reasoning>"), "Copy should accept XML"); assert!(copy.append("<reasoning>"), "Copy should accept XML");
assert!(!processor.eof().unwrap(), "Original should be unchanged"); assert!(!processor.eof().unwrap(), "Original should be unchanged");
assert!(!copy.eof().unwrap(), "Copy should be modified"); 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 "<reasoning>" one by one // Now try to append each character of "<reasoning>" one by one
let chars = vec!['<', 'r', 'e', 'a', 's', 'o', 'n', 'i', 'n', 'g', '>']; let chars = vec!['<', 'r', 'e', 'a', 's', 'o', 'n', 'i', 'n', 'g', '>'];
let mut full_text = String::new(); let mut full_text = String::new();
let processor = processor.copy(py).unwrap(); let processor = processor.copy().unwrap();
for c in chars { for c in chars {
full_text.push(c); full_text.push(c);
@@ -373,7 +374,7 @@ processed_scores = processor(input_ids, scores)
if element_names.contains(&"reasoning".to_string()) { if element_names.contains(&"reasoning".to_string()) {
// Test appending "<reasoning>" to a fresh processor // Test appending "<reasoning>" to a fresh processor
let mut processor = processor.copy(py).unwrap(); let mut processor = processor.copy().unwrap();
let result = processor.append("<reasoning>"); let result = processor.append("<reasoning>");
println!("Appending '<reasoning>' to fresh processor: {}", result); println!("Appending '<reasoning>' to fresh processor: {}", result);

130
lib/xml_schema_validator/tests/test_logits_processor.py
View File

@@ -18,67 +18,67 @@ xml_schema_only_root_node = """<?xml version="1.0" encoding="UTF-8"?>
</xs:element> </xs:element>
</xs:schema>""" </xs:schema>"""
#def test_logits_processor_init_copy(): def test_logits_processor_init_copy():
# """Test the initialization of the LogitsProcessor""" """Test the initialization of the LogitsProcessor"""
#
# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
# logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
# assert logits_processor.core is not None assert logits_processor.core is not None
#
# logits_processor_copy = logits_processor.copy() logits_processor_copy = logits_processor.copy()
# assert logits_processor_copy.core is not None assert logits_processor_copy.core is not None
#
#def test_xml_schema_parsing(): def test_xml_schema_parsing():
# """Test basic XML schema parsing with different element types""" """Test basic XML schema parsing with different element types"""
#
# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
# logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node) logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
# assert logits_processor.core is not None assert logits_processor.core is not None
#
# element_names = logits_processor.core.get_element_names() element_names = logits_processor.core.get_element_names()
# assert len(element_names) > 0 assert len(element_names) > 0
# assert "root" in element_names assert "root" in element_names
#
#def test_basic_xml_validation(): def test_basic_xml_validation():
# """Test basic validation of XML fragments against the schema""" """Test basic validation of XML fragments against the schema"""
#
# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
# processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
#
# # Test appending valid XML fragments # Test appending valid XML fragments
# assert processor.core.copy().append("<reasoning>") assert processor.core.copy().append("<reasoning>")
# assert processor.core.copy().append("<reasoning>Test content</reasoning>") assert processor.core.copy().append("<reasoning>Test content</reasoning>")
#
# # Test appending invalid XML fragments # Test appending invalid XML fragments
# assert not processor.core.copy().append("<invalid>") assert not processor.core.copy().append("<invalid>")
# assert not processor.core.copy().append("<reasoning></invalid>") assert not processor.core.copy().append("<reasoning></invalid>")
#
#def test_token_masking(): def test_token_masking():
# """Test that invalid tokens are properly masked""" """Test that invalid tokens are properly masked"""
#
# tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token) tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
# processor = XmlLogitsProcessor(tokenizer, xml_schema_actions) processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
#
# # Process empty input to set prompt length # Process empty input to set prompt length
# input_ids = torch.tensor([tokenizer.encode("")]) input_ids = torch.tensor([tokenizer.encode("")])
# scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
# processed_scores = processor(input_ids, scores) processed_scores = processor(input_ids, scores)
#
# # Process valid continuation # Process valid continuation
# input_ids = torch.tensor([tokenizer.encode("<reasoning>")]) input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
# scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
# processed_scores = processor(input_ids, scores) processed_scores = processor(input_ids, scores)
#
# # Verify that valid continuations still have positive scores # Verify that valid continuations still have positive scores
# space_token = tokenizer.encode(" ", add_special_tokens=False)[0] space_token = tokenizer.encode(" ", add_special_tokens=False)[0]
# assert processed_scores[0, space_token] > -float('inf') assert processed_scores[0, space_token] > -float('inf')
#
# # Verify that invalid continuations are masked # Verify that invalid continuations are masked
# input_ids = torch.tensor([tokenizer.encode("</invalid>")]) input_ids = torch.tensor([tokenizer.encode("</invalid>")])
# scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability scores = torch.ones((1, len(tokenizer))) # All tokens have equal probability
# processed_scores = processor(input_ids, scores) processed_scores = processor(input_ids, scores)
# assert processed_scores[0, tokenizer.eos_token_id] == 1 assert processed_scores[0, tokenizer.eos_token_id] == 1
# assert processed_scores[0, space_token] == -float('inf') assert processed_scores[0, space_token] == -float('inf')
def test_delete(): def test_delete():
"""Test that invalid tokens are properly masked""" """Test that invalid tokens are properly masked"""
@@ -100,8 +100,8 @@ def test_delete():
if __name__ == "__main__": if __name__ == "__main__":
# Run individual tests for debugging # Run individual tests for debugging
#test_logits_processor_init_copy() test_logits_processor_init_copy()
#test_xml_schema_parsing() test_xml_schema_parsing()
#test_basic_xml_validation() test_basic_xml_validation()
#test_token_masking() test_token_masking()
test_delete() test_delete()