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);