diff --git a/lib/xml_schema_validator/Cargo.toml b/lib/xml_schema_validator/Cargo.toml index 3013479..4cf4733 100644 --- a/lib/xml_schema_validator/Cargo.toml +++ b/lib/xml_schema_validator/Cargo.toml @@ -7,17 +7,14 @@ description = "XML Schema validation library optimized for llm token validation" [lib] name = "xml_schema_validator" -# Required for PyO3 crate-type = ["cdylib", "rlib"] [dependencies] -# XML parsing utilities +numpy = "0.24" +pyo3 = { version = "0.24", features = ["extension-module"] } quick-xml = { version = "0.28.1", features = ["serde", "serialize"] } serde = { version = "1", features = ["derive"] } -thiserror = "1.0.38" # Better error handling - -# Python bindings -pyo3 = { version = "0.18.0", features = ["extension-module"] } +thiserror = "1.0.38" [build-dependencies] xml_schema_generator = { version = "0.6.18", features = ["env_logger"] } \ No newline at end of file diff --git a/lib/xml_schema_validator/src/lib.rs b/lib/xml_schema_validator/src/lib.rs index 8c39441..7af247e 100644 --- a/lib/xml_schema_validator/src/lib.rs +++ b/lib/xml_schema_validator/src/lib.rs @@ -4,6 +4,7 @@ //! A streaming XML validator that checks XML fragments against an XSD schema mod error; +mod logits_processor; mod python; mod schema; mod token; diff --git a/lib/xml_schema_validator/src/logits_processor.rs b/lib/xml_schema_validator/src/logits_processor.rs new file mode 100644 index 0000000..7bf684c --- /dev/null +++ b/lib/xml_schema_validator/src/logits_processor.rs @@ -0,0 +1,174 @@ +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyString}; +use pyo3::exceptions::PyValueError; +use std::sync::RwLock; + +/// A minimal LogitsProcessor that enforces valid XML according to a schema. +#[pyclass] +pub struct XmlLogitsProcessor { + /// The Python tokenizer object + tokenizer: PyObject, + /// The schema text to validate against + schema_text: String, + /// The prompt length (for tracking generated vs. prompt tokens) + prompt_length: RwLock>, + /// Cached EOS token ID + eos_token_id: RwLock>, +} + +#[pymethods] +impl XmlLogitsProcessor { + #[new] + fn new(py: Python<'_>, tokenizer: PyObject, schema_text: &str) -> PyResult { + // Create a small test to verify the schema is valid + match crate::XmlSchemaValidator::new(schema_text) { + Ok(_) => (), + Err(e) => return Err(PyValueError::new_err(format!("Invalid schema: {}", e))), + } + + // Cache the EOS token ID + let eos_token_id = tokenizer.getattr(py, "eos_token_id")?; + let eos_token_id = if eos_token_id.is_none(py) { + None + } else { + Some(eos_token_id.extract::(py)?) + }; + + Ok(Self { + tokenizer, + schema_text: schema_text.to_string(), + prompt_length: RwLock::new(None), + eos_token_id: RwLock::new(eos_token_id), + }) + } + + /// Validate if a text is valid XML according to our schema + fn validate_text(&self, text: &str) -> bool { + // Create a fresh validator for each call + if let Ok(mut validator) = crate::XmlSchemaValidator::new(&self.schema_text) { + validator.append(text).is_ok() + } else { + false + } + } + + /// Process logits to mask invalid XML tokens. + /// This is an unoptimized version that simply processes one token at a time. + fn __call__(&self, py: Python<'_>, input_ids: PyObject, scores: PyObject) -> PyResult { + // Create a simple Python script for processing + let process_code = r#" +import numpy as np + +# Get shapes from the tensors +batch_size, seq_length = input_ids.shape +_, vocab_size = scores.shape + +# If this is the first call, store the prompt length +if prompt_length is None: + prompt_length = seq_length + +result_prompt_length = prompt_length + +# Process each batch item +for batch_idx in range(batch_size): + # Get the input_ids for this batch + current_ids = input_ids[batch_idx].tolist() + + # Extract the generated portion (after the prompt) + generated_ids = current_ids[prompt_length:] if prompt_length < len(current_ids) else [] + + # Decode the generated text + generated_text = tokenizer.decode(generated_ids, skip_special_tokens=False) if generated_ids else "" + + # Check if the current generated text is valid + is_current_valid = processor_obj.validate_text(generated_text) + + # Create a mask for valid tokens + valid_tokens_mask = np.zeros(vocab_size, dtype=bool) + + if not is_current_valid and generated_text: + # If the current text is invalid, only allow EOS token if available + if eos_token_id is not None: + valid_tokens_mask[eos_token_id] = True + else: + # Process each token individually - NO BATCHING FOR SIMPLICITY + for token_idx in range(vocab_size): + # Decode this single token + token_text = tokenizer.decode([token_idx], skip_special_tokens=False) + + # Check if adding this token would produce valid XML + valid = processor_obj.validate_text(generated_text + token_text) + if valid: + valid_tokens_mask[token_idx] = True + + # Apply the mask to scores + invalid_tokens_mask = ~valid_tokens_mask + scores[batch_idx, invalid_tokens_mask] = float('-inf') +"#; + + // Create namespace for execution + let locals = PyDict::new(py); + locals.set_item("input_ids", &input_ids)?; + locals.set_item("scores", &scores)?; + locals.set_item("tokenizer", &self.tokenizer)?; + + // Convert self to a Python object + let processor_obj = PyCell::new(py, self.clone())?; + locals.set_item("processor_obj", processor_obj)?; + + // Handle prompt_length + if let Ok(guard) = self.prompt_length.read() { + match *guard { + Some(pl) => locals.set_item("prompt_length", pl)?, + None => locals.set_item("prompt_length", py.None())?, + } + } else { + locals.set_item("prompt_length", py.None())?; + } + + // Handle eos_token_id + if let Ok(guard) = self.eos_token_id.read() { + match *guard { + Some(eos) => locals.set_item("eos_token_id", eos)?, + None => locals.set_item("eos_token_id", py.None())?, + } + } else { + locals.set_item("eos_token_id", py.None())?; + } + + // Execute the Python code - Using the newer API + py.run(process_code, None, Some(locals))?; + + // Update the prompt length + if let Some(result_pl) = locals.get_item("result_prompt_length") { + if !result_pl.is_none() { + if let Ok(new_pl) = result_pl.extract::() { + if let Ok(mut pl_guard) = self.prompt_length.write() { + *pl_guard = Some(new_pl); + } + } + } + } + + // Return the processed scores + Ok(scores) + } +} + +// Add Clone trait implementation +impl Clone for XmlLogitsProcessor { + fn clone(&self) -> Self { + Self { + tokenizer: self.tokenizer.clone(), + schema_text: self.schema_text.clone(), + prompt_length: RwLock::new(self.prompt_length.read().ok().and_then(|g| *g)), + eos_token_id: RwLock::new(self.eos_token_id.read().ok().and_then(|g| *g)), + } + } +} + +/// Register the XmlLogitsProcessor class with the Python module +pub fn register_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + Ok(()) +} \ No newline at end of file diff --git a/lib/xml_schema_validator/src/python.rs b/lib/xml_schema_validator/src/python.rs index 6654226..6b64435 100644 --- a/lib/xml_schema_validator/src/python.rs +++ b/lib/xml_schema_validator/src/python.rs @@ -33,7 +33,11 @@ impl XmlSchemaValidator { /// Python module for XML Schema validation #[pymodule] -fn xml_schema_validator(_py: Python<'_>, m: &PyModule) -> PyResult<()> { +fn xml_schema_validator(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; + + // Register the XmlLogitsProcessor class + crate::logits_processor::register_module(py, m)?; + Ok(()) } \ No newline at end of file