wip logits processor in ruuuuuuust

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"] }

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;

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<Option<usize>>,
+    /// Cached EOS token ID
+    eos_token_id: RwLock<Option<i64>>,
+}
+
+#[pymethods]
+impl XmlLogitsProcessor {
+    #[new]
+    fn new(py: Python<'_>, tokenizer: PyObject, schema_text: &str) -> PyResult<Self> {
+        // 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::<i64>(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<PyObject> {
+        // 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::<usize>() {
+                    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::<XmlLogitsProcessor>()?;
+    Ok(())
+}

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::<XmlSchemaValidator>()?;
+    
+    // Register the XmlLogitsProcessor class
+    crate::logits_processor::register_module(py, m)?;
+    
     Ok(())
 }