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

Integrate logits processor with qwq (untested)

Browse Source
This commit is contained in:
Niels Geens
2025-04-07 13:35:20 +02:00
parent 023fc75aac
commit 50539f6d03
5 changed files with 101 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
lib/xml_schema_validator/pyproject.toml
View File

@@ -5,19 +5,9 @@ build-backend = "maturin"
[project]
name = "xml_schema_validator"
version = "0.1.0"
description = "XML Schema validation library for SIA"
description = "XML Schema validation library"
requires-python = ">=3.10"
classifiers = [
"Programming Language :: Rust",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Topic :: Text Processing :: Markup :: XML",
]
dependencies = []
[project.urls]
repository = "https://github.com/sia/xml_schema_validator"
[tool.maturin]
features = ["pyo3/extension-module"]

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

@@ -13,7 +13,7 @@ use token::*;
/// An XML validator that checks XML fragments against an XSD schema
#[derive(Clone, Debug)]
pub struct XmlValidator {
pub struct XmlSchemaValidator {
current_tokens: Vec<Box<Token>>,
}
@@ -22,7 +22,7 @@ pub struct XmlValidator {
/// This is a type alias for `std::result::Result` with our custom `Error` type
pub type Result<T> = std::result::Result<T, Error>;
impl XmlValidator {
impl XmlSchemaValidator {
/// Create a new validator with the given schema text
pub fn new(schema_text: &str) -> Result<Self> {
let schema: schema::XsSchema = quick_xml::de::from_str(schema_text)?;
@@ -66,7 +66,7 @@ mod tests {
fn test_element_open() {
let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let input = "<reasoning>";
let mut validator = XmlValidator::new(&schema_text).unwrap();
let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
validator.append(input).unwrap();
}
@@ -74,7 +74,7 @@ mod tests {
fn test_element_open_whitespace() {
let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let input = "<reasoning >";
let mut validator = XmlValidator::new(&schema_text).unwrap();
let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
validator.append(input).unwrap();
}
@@ -82,7 +82,7 @@ mod tests {
fn test_element_open_invalid_name() {
let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let input = "<rae";
let mut validator = XmlValidator::new(&schema_text).unwrap();
let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
validator.append(input).expect_err("Should fail");
}
@@ -90,7 +90,7 @@ mod tests {
fn test_element_close() {
let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let input = "<reasoning>test</reasoning>";
let mut validator = XmlValidator::new(&schema_text).unwrap();
let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
validator.append(input).unwrap();
}
@@ -98,7 +98,7 @@ mod tests {
fn test_element_close_mismatched() {
let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
let input = "<reasoning>test</wrong>";
let mut validator = XmlValidator::new(&schema_text).unwrap();
let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
validator.append(input).expect_err("Should fail with mismatched closing tag");
}
}

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

@@ -1,18 +1,17 @@
use crate::XmlValidator;
use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;
/// Python wrapper for XmlValidator
#[pyclass(unsendable)]
struct PyXmlValidator {
validator: XmlValidator,
struct XmlSchemaValidator {
validator: crate::XmlSchemaValidator,
}
#[pymethods]
impl PyXmlValidator {
impl XmlSchemaValidator {
#[new]
fn new(schema_text: &str) -> PyResult<Self> {
let validator = XmlValidator::new(schema_text)
let validator = crate::XmlSchemaValidator::new(schema_text)
.map_err(|e| PyValueError::new_err(format!("Failed to create validator: {}", e)))?;
Ok(Self { validator })
@@ -22,11 +21,17 @@ impl PyXmlValidator {
self.validator.append(fragment)
.map_err(|e| PyValueError::new_err(format!("Validation error: {}", e)))
}
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: &PyModule) -> PyResult<()> {
m.add_class::<PyXmlValidator>()?;
m.add_class::<XmlSchemaValidator>()?;
Ok(())
}

28
sia/llm_engine/qwq_llm_engine.py
View File

@@ -1,12 +1,13 @@
from pathlib import Path
from threading import Thread
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
from typing import Callable, Iterator
from typing import Callable, Iterator, Optional, List
import json
import torch
from . import LlmEngine
from .. import util
from .xml_logits_processor import XmlLogitsProcessor
class QwQLlmEngine(LlmEngine):
@@ -15,6 +16,7 @@ class QwQLlmEngine(LlmEngine):
model_path: Path,
temperature: float,
token_limit: int = None,
xml_schema_text: Optional[str] = None,
):
"""
Initialize the QwQ LLM Engine.
@@ -23,9 +25,12 @@ class QwQLlmEngine(LlmEngine):
model_path: Local path to the model
temperature: Sampling temperature
token_limit: Maximum tokens to generate
xml_schema_text: Optional XML schema to validate against
"""
self._temperature = temperature
self._token_limit = token_limit
if xml_schema_text:
self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
tokenizer_config = json.load(f)
@@ -86,17 +91,22 @@ class QwQLlmEngine(LlmEngine):
skip_prompt=True,
)
generation_kwargs = {
"text_inputs": text,
"do_sample": True,
"temperature": self._temperature,
"max_new_tokens": self._token_limit,
"streamer": streamer,
}
if self._logits_processor:
generation_kwargs["logits_processor"] = self.logits_processor
generation_thread = Thread(
target=self._pipline,
kwargs=dict(
text_inputs=text,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self._token_limit,
streamer=streamer,
)
kwargs=generation_kwargs
)
generation_thread.start()
for text in util.stop_before_value(streamer, self._tokenizer.eos_token):

63
sia/llm_engine/xml_logits_processor.py Normal file
View File

@@ -0,0 +1,63 @@
import torch
from transformers import LogitsProcessor
from xml_schema_validator import XmlSchemaValidator
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, 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.tokenizer = tokenizer
self.schema_validator = XmlSchemaValidator(schema_text)
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
# For each sequence in the batch
for batch_idx in range(batch_size):
# Get the current text generated so far
current_ids = input_ids[batch_idx]
current_text = self.tokenizer.decode(current_ids)
# Get all possible next tokens
vocab_size = scores.shape[-1]
# Create a mask to track which tokens are valid
valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
# For each possible next token
for token_idx in range(vocab_size):
# Create a copy of the validator to test this token
validator_copy = self.schema_validator.copy()
# Decode the token and test if appending it would be valid
token_text = self.tokenizer.decode([token_idx])
if validator_copy.append(token_text):
valid_tokens_mask[token_idx] = True
# Mask out invalid tokens by setting their scores to negative infinity
invalid_tokens_mask = ~valid_tokens_mask
scores[batch_idx, invalid_tokens_mask] = float('-inf')
return scores