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

Added char trie

Browse Source
This commit is contained in:
Niels Geens
2025-04-10 21:46:24 +02:00
parent 10df42c3a6
commit ba9d69b414
14 changed files with 126 additions and 730 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
lib/xml_schema_validator/src/char_trie.rs Normal file
View File

@@ -0,0 +1,103 @@
use crate::*;
use pyo3::prelude::*;
// Node in our character trie
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>>,
}
impl CharTrieNode {
fn new() -> Self {
Self {
children: Vec::new(),
token_id: None,
}
}
// Add a child for character c, maintaining sorted order
fn add_child(&mut self, c: char) -> &mut CharTrieNode {
match self.children.binary_search_by_key(&c, |(ch, _)| *ch) {
Ok(index) => &mut self.children[index].1,
Err(index) => {
// Insert at the right position to maintain sort order
self.children.insert(index, (c, CharTrieNode::new()));
&mut self.children[index].1
}
}
}
// Add a token id to this node
fn add_token(&mut self, token_id: Py<PyAny>) {
self.token_id = Some(token_id);
}
pub 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
pub struct CharTrie {
root: CharTrieNode,
}
impl CharTrie {
pub fn new() -> Self {
Self {
root: CharTrieNode::new(),
}
}
// Insert a token's character sequence into the trie
pub fn insert(&mut self, token_text: &str, token_id: Py<PyAny>) {
let mut current = &mut self.root;
// Add each character in the token text
for c in token_text.chars() {
current = current.add_child(c);
}
// Mark this node as representing token_id
current.add_token(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, token_map: &Vec<(Py<PyAny>, String)>) -> Vec<Py<PyAny>> {
let mut invalid_tokens = Vec::new();
for (token_id, token_text) in token_map {
// Check if this token would make the XML invalid
if !self.is_valid_continuation(xml_validator, token_text) {
invalid_tokens.push(token_id.clone_ref(py));
}
}
invalid_tokens
}
// Check if a token would be a valid continuation
fn is_valid_continuation(&self, xml_validator: &XmlSchemaValidator, token_text: &str) -> bool {
// First, quick check with just the first character
if let Some(first_char) = token_text.chars().next() {
if !xml_validator.can_continue_with(first_char) {
return false;
}
}
// If first character is valid, try the full token
let mut validator_clone = xml_validator.clone();
validator_clone.append(token_text).is_ok()
}
pub fn clone_ref(&self, py: Python<'_>) -> Self {
Self {
root: self.root.clone_ref(py),
}
}
}

3
lib/xml_schema_validator/src/lib.rs
View File

@@ -3,10 +3,11 @@
//! A streaming XML validator that checks XML fragments against an XSD schema
mod char_trie;
mod error;
mod python;
mod py_xml_logits_processor_core;
mod py_xml_schema_validator;
mod python;
mod schema;
mod token;

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

@@ -1,25 +1,32 @@
use pyo3::prelude::*;
use pyo3::types::PyDict;
use crate::*;
/// A minimal LogitsProcessor that enforces valid XML according to a schema.
#[pyclass]
pub struct XmlLogitsProcessorCore {
xml_schema_validator: crate::XmlSchemaValidator,
xml_schema_validator: XmlSchemaValidator,
tokens: Vec<(PyObject, String)>,
trie: char_trie::CharTrie,
}
#[pymethods]
impl XmlLogitsProcessorCore {
#[new]
fn new(tokens: Bound<'_, PyDict>, schema_text: &str) -> PyResult<Self> {
let xml_schema_validator = match crate::XmlSchemaValidator::new(schema_text) {
fn new(py: Python<'_>, tokens: Bound<'_, PyDict>, schema_text: &str) -> PyResult<Self> {
let xml_schema_validator = match XmlSchemaValidator::new(schema_text) {
Ok(xml_schema_validator) => xml_schema_validator,
Err(e) => return Err(pyo3::exceptions::PyValueError::new_err(e.to_string())),
};
let tokens = tokens.iter().map(|(id, value)| (id.unbind(), value.unbind().to_string())).collect();
let tokens: Vec<(Py<PyAny>, String)> = tokens.iter().map(|(id, value)| (id.unbind(), value.unbind().to_string())).collect();
let mut trie = char_trie::CharTrie::new();
for (token_id, token_text) in tokens.iter() {
trie.insert(token_text, token_id.clone_ref(py));
}
Ok(Self {
xml_schema_validator,
tokens,
trie,
})
}
@@ -31,20 +38,7 @@ impl XmlLogitsProcessorCore {
/// Get a list of tokens that would lead to invalid XML
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()
.filter(|token| !chars[token.1.chars().next().unwrap() as usize])
.filter_map(|token| {
if self.xml_schema_validator.clone().append(&token.1).is_ok() {
None
} else {
Some(token.0.clone_ref(py))
}
})
.collect::<Vec<_>>())
Ok(self.trie.find_invalid_tokens(py, &self.xml_schema_validator, &self.tokens))
}
/// Check if the validator has reached the end of the XML
@@ -56,7 +50,8 @@ impl XmlLogitsProcessorCore {
fn copy(&self, py: Python<'_>) -> PyResult<Self> {
Ok(Self {
xml_schema_validator: self.xml_schema_validator.clone(),
tokens: self.tokens.iter().map(|(a, b)|(a.clone_ref(py), b.clone())).collect()
tokens: self.tokens.iter().map(|(a, b)|(a.clone_ref(py), b.clone())).collect(),
trie: self.trie.clone_ref(py),
})
}
@@ -109,7 +104,7 @@ mod tests {
</xs:element>
</xs:schema>"#;
XmlLogitsProcessorCore::new(tokens.into(), schema_text)
XmlLogitsProcessorCore::new(py, tokens.into(), schema_text)
}
#[test]
@@ -131,7 +126,7 @@ mod tests {
tokens.set_item(1, "reasoning").unwrap();
tokens.set_item(2, ">").unwrap();
let mut processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
let mut processor = XmlLogitsProcessorCore::new(py, tokens.into(), &schema_text).unwrap();
assert!(processor.append("<reasoning>"),
"Should accept <reasoning> with actual schema");
}
@@ -233,7 +228,7 @@ mod tests {
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(py, tokens.into(), &schema_text).unwrap();
let names = processor.get_element_names().unwrap();
// Print out all elements from schema for debugging
@@ -275,7 +270,7 @@ mod tests {
tokens.set_item(10, ">").unwrap();
// Create processor with this schema and tokens
let processor = XmlLogitsProcessorCore::new(tokens.clone().into(), schema_text).unwrap();
let processor = XmlLogitsProcessorCore::new(py, tokens.clone().into(), schema_text).unwrap();
// Test basic operations to verify core functionality works
let element_names = processor.get_element_names().unwrap();
@@ -296,7 +291,7 @@ mod tests {
c, &full_text[..full_text.len()-1], result);
// Create a new processor to test appending the fragment built so far
let mut fresh_processor = XmlLogitsProcessorCore::new(tokens.clone().into(), schema_text).unwrap();
let mut fresh_processor = XmlLogitsProcessorCore::new(py, tokens.clone().into(), schema_text).unwrap();
let append_result = fresh_processor.append(&full_text);
println!("Appending full text '{}': {}", full_text, append_result);
}
@@ -305,7 +300,7 @@ mod tests {
if let Ok(actual_schema) = std::fs::read_to_string("../../action_schema.xsd") {
println!("\nTesting with actual schema from disk:");
let processor = XmlLogitsProcessorCore::new(tokens.clone().into(), &actual_schema).unwrap();
let processor = XmlLogitsProcessorCore::new(py, tokens.clone().into(), &actual_schema).unwrap();
let element_names = processor.get_element_names().unwrap();
println!("Element names in actual schema: {:?}", element_names);
@@ -321,7 +316,7 @@ mod tests {
let mut text = String::new();
for c in "<reasoning>".chars() {
text.push(c);
let result = XmlLogitsProcessorCore::new(tokens.clone().into(), &actual_schema).unwrap()
let result = XmlLogitsProcessorCore::new(py, tokens.clone().into(), &actual_schema).unwrap()
.append(&text);
println!("Appending '{}' to fresh processor: {}", text, result);
}