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

2
lib/xml_schema_validator/.dockerignore
View File

@@ -1,2 +0,0 @@
target/
Cargo.lock

6
lib/xml_schema_validator/Cargo.toml
View File

@@ -10,11 +10,7 @@ name = "xml_schema_validator"
crate-type = ["cdylib", "rlib"]
[dependencies]
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"
[build-dependencies]
xml_schema_generator = { version = "0.6.18", features = ["env_logger"] }
thiserror = "1.0.38"

13
lib/xml_schema_validator/pyproject.toml
View File

@@ -1,13 +1,16 @@
[build-system]
requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[project]
name = "xml_schema_validator"
version = "0.1.0"
description = "XML Schema validation library"
requires-python = ">=3.10"
dependencies = []
dependencies = [
"pytest>=8",
"maturin>=1.0,<2.0",
]
[build-system]
requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[tool.maturin]
features = ["pyo3/extension-module"]

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(),
))]
}
}

3
lib/xml_schema_validator/test_requirements.txt Normal file
View File

@@ -0,0 +1,3 @@
torch>=2.0.0
transformers>=4.30.0
.

9
lib/xml_schema_validator/tests/conftest.py Normal file
View File

@@ -0,0 +1,9 @@
import subprocess
import sys
def pytest_sessionstart(session):
"""
Build the Rust extension module before running tests.
"""
print("Installing test dependencies...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "test_requirements.txt"])

71
lib/xml_schema_validator/tests/test_logits_processor.py Normal file
View File

@@ -0,0 +1,71 @@
from xml_logits_processor import XmlLogitsProcessor
from threading import Thread
from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
import os
import sys
def test_logits_processor():
print("test_logits_processor")
model_name = "../../model"
xml_schema_text = """<?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>"""
model = "google/gemma-3-1b-it"
pipline = pipeline(
"text-generation",
model=model,
token=os.environ["SIA_HF_API_KEY"],
)
tokenizer = AutoTokenizer.from_pretrained(
model,
token=os.environ["SIA_HF_API_KEY"],
)
logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)
messages = [
{"role": "system", "content": "Always respond with <root>message by the user</root>. So if the user says 'hello world', you response <root>hello world</root>"},
{"role": "user", "content": "hello, I am the user"}
]
text_inputs = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True
)
generation_kwargs = {
"text_inputs": text_inputs,
"max_new_tokens": 20,
"streamer": streamer,
#"logits_processor": [logits_processor],
}
generation_thread = Thread(
target=pipline,
kwargs=generation_kwargs
)
generation_thread.start()
for text in streamer:
print(text, end="", file=sys.stderr, flush=True)
generation_thread.join()
if __name__ == "__main__":
test_logits_processor()

75
lib/xml_schema_validator/xml_schema_validator/__init__.py Normal file
View File

@@ -0,0 +1,75 @@
import torch
import sys
from transformers import LogitsProcessor
from transformers import AutoTokenizer
import xml_schema_validator
class XmlLogitsProcessor(LogitsProcessor):
"""
A LogitsProcessor that enforces valid XML according to a schema.
This processor masks tokens that would lead to invalid XML
by setting their logits to negative infinity.
"""
def __init__(self, tokenizer: AutoTokenizer, schema_text: str):
"""
Initialize the processor with a schema and tokenizer.
Args:
tokenizer: The tokenizer to use for decoding tokens
schema_text: The XSD schema text to validate against
"""
self.eos_token_id = tokenizer.eos_token_id
self.tokenizer = tokenizer
vocab = tokenizer.get_vocab() # This is {token: id}
items = dict()
for token, id in vocab.items():
items[id] = token
self.core = xml_schema_validator.XmlLogitsProcessorCore(items, schema_text)
self.prompt_length = None
self.is_first_call = True
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
"""
Process logits to mask invalid XML tokens.
Args:
input_ids: Current input token IDs
scores: Current scores (logits) for next token prediction
Returns:
Processed scores with invalid tokens masked
"""
batch_size, _ = input_ids.shape
# If this is the first call, store the prompt length
if self.is_first_call:
self.prompt_length = input_ids.shape[1]
self.is_first_call = False
# For each sequence in the batch
for batch_idx in range(batch_size):
# Get only the generated portion of the text
current_ids = input_ids[batch_idx]
generated_ids = current_ids[self.prompt_length:]
generated_text = self.tokenizer.decode(generated_ids)
# Create a mask to track which tokens are valid
batch_processor = self.core.copy()
if generated_text:
batch_processor.append(generated_text)
if batch_processor.eof():
vocab_size = scores.shape[-1]
valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
valid_tokens_mask[self.eos_token_id] = True
invalid_tokens_mask = ~valid_tokens_mask
scores[batch_idx, invalid_tokens_mask] = float('-inf')
continue
for token in batch_processor.get_invalid_tokens():
if token < scores.shape[1]:
scores[batch_idx, token] = float('-inf')
return scores