Integrate logits processor with qwq (untested)

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

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<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");
     }
 }

lib/xml_schema_validator/src/python.rs

@@ -1,18 +1,17 @@
-use crate::XmlValidator;
 use pyo3::prelude::*;
 use pyo3::exceptions::PyValueError;
 
 /// Python wrapper for XmlValidator
 #[pyclass(unsendable)]
-struct PyXmlValidator {
+struct XmlSchemaValidator {
-    validator: XmlValidator,
+    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(())
 }

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,15 +91,20 @@ 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(
+            kwargs=generation_kwargs
-                text_inputs=text,
-                do_sample=True,
-                temperature=self._temperature,
-                max_new_tokens=self._token_limit,
-                streamer=streamer,
-            )
         )
 
         generation_thread.start()

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