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

Added small optimization to logits processor core

Browse Source
This commit is contained in:
Niels Geens
2025-04-10 19:11:59 +02:00
parent d81ae09f67
commit 10df42c3a6
1 changed files with 18 additions and 136 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,7 +31,12 @@ impl XmlLogitsProcessorCore {
/// Get a list of tokens that would lead to invalid XML /// Get a list of tokens that would lead to invalid XML
fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> { fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
let mut chars = [false; 127];
for i in 0..chars.len() {
chars[i] = self.xml_schema_validator.append_char(i as u8 as char).is_ok();
}
Ok(self.tokens.iter() Ok(self.tokens.iter()
.filter(|token| !chars[token.1.chars().next().unwrap() as usize])
.filter_map(|token| { .filter_map(|token| {
if self.xml_schema_validator.clone().append(&token.1).is_ok() { if self.xml_schema_validator.clone().append(&token.1).is_ok() {
None None
@@ -41,19 +46,6 @@ impl XmlLogitsProcessorCore {
}) })
.collect::<Vec<_>>()) .collect::<Vec<_>>())
} }
/// Get a list of tokens that would lead to valid XML
fn get_valid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
Ok(self.tokens.iter()
.filter_map(|token| {
if self.xml_schema_validator.clone().append(&token.1).is_ok() {
Some(token.0.clone_ref(py))
} else {
None
}
})
.collect::<Vec<_>>())
}
/// Check if the validator has reached the end of the XML /// Check if the validator has reached the end of the XML
fn eof(&self) -> PyResult<bool> { fn eof(&self) -> PyResult<bool> {
@@ -81,31 +73,6 @@ impl XmlLogitsProcessorCore {
Ok(false) Ok(false)
} }
} }
/// Reset the validator to its initial state
fn reset(&mut self) -> PyResult<()> {
// Create a new validator with the same schema elements
let schema_text = r#"<?xml version="1.0" encoding="UTF-8"?>
<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
<xs:element name="root">
<xs:complexType mixed="true">
<xs:sequence>
<xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
</xs:sequence>
</xs:complexType>
</xs:element>
</xs:schema>"#;
match crate::XmlSchemaValidator::new(schema_text) {
Ok(validator) => {
self.xml_schema_validator = validator;
Ok(())
},
Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
"Failed to reset validator: {}", e
))),
}
}
} }
/// Register the class with the Python module /// Register the class with the Python module
@@ -185,43 +152,6 @@ mod tests {
}); });
} }
#[test]
fn test_get_valid_tokens() {
Python::with_gil(|py| {
let mut processor = create_test_processor(py).unwrap();
// At start, only "<" is valid
let valid_tokens = processor.get_valid_tokens(py).unwrap();
let contains_0 = valid_tokens.iter().any(|obj| {
let obj_id: Option<i64> = obj.extract(py).ok();
obj_id == Some(0)
});
assert!(contains_0, "Token '<' should be valid at start");
assert_eq!(valid_tokens.len(), 1, "Only one token should be valid at start");
// After <reasoning>, any character should be valid in mixed content
let mut processor = create_test_processor(py).unwrap();
assert!(processor.append("<reasoning>"), "Should be able to append <reasoning>");
let valid_tokens = processor.get_valid_tokens(py).unwrap();
let contains_3 = valid_tokens.iter().any(|obj| {
let obj_id: Option<i64> = obj.extract(py).ok();
obj_id == Some(3)
});
let contains_4 = valid_tokens.iter().any(|obj| {
let obj_id: Option<i64> = obj.extract(py).ok();
obj_id == Some(4)
});
assert!(contains_3, "Token ' ' should be valid inside mixed content");
assert!(contains_4, "Token 'a' should be valid inside mixed content");
});
}
#[test] #[test]
fn test_eof_detection() { fn test_eof_detection() {
Python::with_gil(|py| { Python::with_gil(|py| {
@@ -299,21 +229,19 @@ mod tests {
"Should contain 'reasoning' element"); "Should contain 'reasoning' element");
assert_eq!(names.len(), 1, "Should have exactly one element"); assert_eq!(names.len(), 1, "Should have exactly one element");
// Test with real schema from disk if available let tokens = PyDict::new(py);
if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") { tokens.set_item(0, "<").unwrap();
let tokens = PyDict::new(py);
tokens.set_item(0, "<").unwrap(); let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
let processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap(); let names = processor.get_element_names().unwrap();
let names = processor.get_element_names().unwrap();
// Print out all elements from schema for debugging
// Print out all elements from schema for debugging println!("Elements in schema: {:?}", names);
println!("Elements in schema: {:?}", names);
// Check if reasoning exists
// Check if reasoning exists assert!(names.contains(&"reasoning".to_string()),
assert!(names.contains(&"reasoning".to_string()), "Real schema should contain 'reasoning'");
"Real schema should contain 'reasoning'");
}
}); });
} }
@@ -403,50 +331,4 @@ mod tests {
} }
}); });
} }
#[test]
fn test_debug_token_processing() {
Python::with_gil(|py| {
if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") {
// Create a dictionary mapping token IDs to their string values
// Simulating how the tokenizer vocabulary works
let tokens = PyDict::new(py);
tokens.set_item(1000, "<").unwrap();
tokens.set_item(1001, "reasoning").unwrap();
tokens.set_item(1002, ">").unwrap();
tokens.set_item(1003, " ").unwrap();
tokens.set_item(1004, "test").unwrap();
tokens.set_item(1005, "</").unwrap();
tokens.set_item(1006, "</reasoning>").unwrap();
let mut processor = XmlLogitsProcessorCore::new(tokens.clone().into(), &schema_text).unwrap();
// Test which tokens are valid at the start
let invalid_tokens = processor.get_invalid_tokens(py).unwrap();
let valid_tokens = processor.get_valid_tokens(py).unwrap();
println!("At start - valid tokens count: {}", valid_tokens.len());
println!("At start - invalid tokens count: {}", invalid_tokens.len());
// Create a new processor and progress to after "<reasoning>"
let mut processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
let result = processor.append("<reasoning>");
println!("Append '<reasoning>': {}", result);
if result {
// Now get valid/invalid tokens after "<reasoning>"
let invalid_tokens = processor.get_invalid_tokens(py).unwrap();
let valid_tokens = processor.get_valid_tokens(py).unwrap();
println!("After '<reasoning>' - valid tokens count: {}", valid_tokens.len());
println!("After '<reasoning>' - invalid tokens count: {}", invalid_tokens.len());
// At this point we'd expect to have more valid tokens than at the start
// since mixed content allows any text
}
} else {
println!("Warning: Couldn't find schema file for testing");
}
});
}
} }