diff --git a/lib/xml_schema_validator/pyproject.toml b/lib/xml_schema_validator/pyproject.toml index 06b71dc..7587057 100644 --- a/lib/xml_schema_validator/pyproject.toml +++ b/lib/xml_schema_validator/pyproject.toml @@ -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"] diff --git a/lib/xml_schema_validator/src/lib.rs b/lib/xml_schema_validator/src/lib.rs index bac12aa..ea23073 100644 --- a/lib/xml_schema_validator/src/lib.rs +++ b/lib/xml_schema_validator/src/lib.rs @@ -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>, } @@ -22,7 +22,7 @@ pub struct XmlValidator { /// This is a type alias for `std::result::Result` with our custom `Error` type pub type Result = std::result::Result; -impl XmlValidator { +impl XmlSchemaValidator { /// Create a new validator with the given schema text pub fn new(schema_text: &str) -> Result { 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 = ""; - 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 = ""; - 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 = " PyResult { - 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 { + 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::()?; + m.add_class::()?; Ok(()) } \ No newline at end of file diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index e4e6e02..94740be 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/sia/llm_engine/qwq_llm_engine.py @@ -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): diff --git a/sia/llm_engine/xml_logits_processor.py b/sia/llm_engine/xml_logits_processor.py new file mode 100644 index 0000000..b7b8553 --- /dev/null +++ b/sia/llm_engine/xml_logits_processor.py @@ -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 \ No newline at end of file