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

Optimized hot-loop for logits processor

Browse Source
This commit is contained in:
Niels Geens
2025-04-09 11:59:26 +02:00
parent 654c32c7ac
commit 8fb30c8ed0
17 changed files with 239 additions and 272 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,8 +4,9 @@
//! A streaming XML validator that checks XML fragments against an XSD schema
mod error;
mod logits_processor;
mod python;
mod py_xml_logits_processor_core;
mod py_xml_schema_validator;
mod schema;
mod token;
@@ -48,6 +49,11 @@ impl XmlSchemaValidator {
Ok(())
}
/// Check if the validator has reached the end of the XML
pub fn eof(&self) -> bool {
self.current_tokens.is_empty() || self.current_tokens.iter().all(Token::is_eof)
}
/// Append a single character to the validator
fn append_char(&mut self, c: char) -> Result<()> {
let mut new_tokens = vec![];

174
lib/xml_schema_validator/src/logits_processor.rs
View File

@@ -1,174 +0,0 @@
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(())
}

58
lib/xml_schema_validator/src/py_xml_logits_processor_core.rs Normal file
View File

@@ -0,0 +1,58 @@
use pyo3::prelude::*;
use pyo3::types::PyDict;
/// A minimal LogitsProcessor that enforces valid XML according to a schema.
#[pyclass]
pub struct XmlLogitsProcessorCore {
xml_schema_validator: crate::XmlSchemaValidator,
tokens: Vec<(PyObject, String)>,
}
#[pymethods]
impl XmlLogitsProcessorCore {
#[new]
fn new(tokens: Bound<'_, PyDict>, schema_text: &str) -> PyResult<Self> {
let xml_schema_validator = match crate::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();
Ok(Self {
xml_schema_validator,
tokens,
})
}
fn append(&mut self, fragment: &str) -> bool {
self.xml_schema_validator.append(fragment).is_ok()
}
fn get_invalid_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() {
None
} else {
Some(token.0.clone_ref(py))
}
})
.collect::<Vec<_>>())
}
fn eof(&self) -> PyResult<bool> {
Ok(self.xml_schema_validator.eof())
}
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()
})
}
}
/// Register the class with the Python module
pub fn register_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<XmlLogitsProcessorCore>()?;
Ok(())
}

42
lib/xml_schema_validator/src/py_xml_schema_validator.rs Normal file
View File

@@ -0,0 +1,42 @@
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// Python wrapper for XmlValidator
#[pyclass]
struct XmlSchemaValidator {
validator: crate::XmlSchemaValidator,
}
#[pymethods]
impl XmlSchemaValidator {
#[new]
fn new(schema_text: &str) -> PyResult<Self> {
let validator = crate::XmlSchemaValidator::new(schema_text)
.map_err(|e| PyValueError::new_err(format!("Failed to create validator: {}", e)))?;
Ok(Self { validator })
}
fn append(&mut self, fragment: &str) -> PyResult<(bool, Option<String>)> {
match self.validator.append(fragment) {
Ok(()) => PyResult::Ok((true, None)),
Err(e) => PyResult::Ok((false, Some(e.to_string()))),
}
}
fn eof(&self) -> PyResult<bool> {
Ok(self.validator.eof())
}
fn copy(&self) -> PyResult<Self> {
Ok(Self {
validator: self.validator.clone(),
})
}
}
/// Register the class with the Python module
pub fn register_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<XmlSchemaValidator>()?;
Ok(())
}

38
lib/xml_schema_validator/src/python.rs
View File

@@ -1,43 +1,9 @@
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// Python wrapper for XmlValidator
#[pyclass]
struct XmlSchemaValidator {
validator: crate::XmlSchemaValidator,
}
#[pymethods]
impl XmlSchemaValidator {
#[new]
fn new(schema_text: &str) -> PyResult<Self> {
let validator = crate::XmlSchemaValidator::new(schema_text)
.map_err(|e| PyValueError::new_err(format!("Failed to create validator: {}", e)))?;
Ok(Self { validator })
}
fn append(&mut self, fragment: &str) -> PyResult<(bool, Option<String>)> {
match self.validator.append(fragment) {
Ok(()) => PyResult::Ok((true, None)),
Err(e) => PyResult::Ok((false, Some(e.to_string()))),
}
}
fn copy(&self) -> PyResult<Self> {
Ok(Self {
validator: self.validator.clone(),
})
}
}
/// Python module for XML Schema validation
#[pymodule]
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)?;
crate::py_xml_logits_processor_core::register_module(py, m)?;
crate::py_xml_schema_validator::register_module(py, m)?;
Ok(())
}

7
lib/xml_schema_validator/src/token/mod.rs
View File

@@ -50,4 +50,11 @@ impl Token {
Token::Whitespace(whitespace) => whitespace.append(c),
}
}
pub fn is_eof(&self) -> bool {
match self {
Token::EndOfFile(_) => true,
_ => false,
}
}
}

13
lib/xml_schema_validator/src/token/text_content.rs
View File

@@ -5,22 +5,12 @@ use crate::*;
#[derive(Clone, Debug)]
pub struct TextContent {
element: Arc<schema::XsElement>,
buffer: String,
}
impl TextContent {
pub fn new(element: Arc<schema::XsElement>) -> Self {
Self {
element,
buffer: String::new(),
}
}
/// Create a new TextContent with the given initial buffer
pub fn with_buffer(element: Arc<schema::XsElement>, buffer: String) -> Self {
Self {
element,
buffer,
}
}
@@ -28,9 +18,8 @@ impl TextContent {
if c == '<' {
vec![Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&self.element)))]
} else {
vec![Token::TextContent(TextContent::with_buffer(
vec![Token::TextContent(TextContent::new(
Arc::clone(&self.element),
self.buffer + &c.to_string(),
))]
}
}