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::*;
// 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<Py<PyAny>>,
token_id: Option<Arc<Py<PyAny>>>,
}
impl CharTrieNode {
@@ -25,15 +25,14 @@ impl CharTrieNode {
}
// 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);
}
fn find_invalid_tokens(
&self,
py: Python<'_>,
xml_validator: &XmlSchemaValidator,
) -> Vec<Py<PyAny>> {
) -> Vec<Arc<Py<PyAny>>> {
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<Py<PyAny>> {
// 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<Arc<Py<PyAny>>> {
// 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<Py<PyAny>> {
fn all_tokens(&self) -> Vec<Arc<Py<PyAny>>> {
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<Py<PyAny>> {
) -> Vec<Arc<Py<PyAny>>> {
// 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)
}
}

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.
#[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<Vec<PyObject>> {
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<Self> {
Ok(Self {
xml_schema_validator: self.xml_schema_validator.clone(),
trie: self.trie.clone_ref(py),
})
fn copy(&self) -> PyResult<Self> {
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("<reasoning>"), "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 "<reasoning>" 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 "<reasoning>" to a fresh processor
let mut processor = processor.copy(py).unwrap();
let mut processor = processor.copy().unwrap();
let result = processor.append("<reasoning>");
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:schema>"""
#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("<reasoning>")
# assert processor.core.copy().append("<reasoning>Test content</reasoning>")
#
# # Test appending invalid XML fragments
# assert not processor.core.copy().append("<invalid>")
# assert not processor.core.copy().append("<reasoning></invalid>")
#
#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("<reasoning>")])
# 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("</invalid>")])
# 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("<reasoning>")
assert processor.core.copy().append("<reasoning>Test content</reasoning>")
# Test appending invalid XML fragments
assert not processor.core.copy().append("<invalid>")
assert not processor.core.copy().append("<reasoning></invalid>")
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("<reasoning>")])
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("</invalid>")])
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()