Optimized hot-loop for logits processor

.dockerignore

@@ -1,2 +1,4 @@
-.env
+.env
-./model/
+./model/
+**/target/
+**/Cargo.lock

Dockerfile

@@ -89,4 +89,4 @@ RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile
 WORKDIR /root/desktop
 
 ENTRYPOINT ["/bin/bash", "-c"]
-CMD ["/root/sia/scripts/install.sh; /bin/bash -c /root/sia/scripts/restart.sh"]
+CMD ["/root/sia/scripts/install.sh; /bin/bash -lc /root/sia/scripts/restart.sh"]

lib/xml_schema_validator/.dockerignore

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

lib/xml_schema_validator/Cargo.toml

@@ -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"
+thiserror = "1.0.38"
-
-[build-dependencies]
-xml_schema_generator = { version = "0.6.18", features = ["env_logger"] }

lib/xml_schema_validator/pyproject.toml

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

lib/xml_schema_validator/src/lib.rs

@@ -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![];

lib/xml_schema_validator/src/logits_processor.rs

@@ -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(())
-}

lib/xml_schema_validator/src/py_xml_logits_processor_core.rs

@@ -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(())
+}

lib/xml_schema_validator/src/py_xml_schema_validator.rs

@@ -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(())
+}

lib/xml_schema_validator/src/python.rs

@@ -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>()?;
+    crate::py_xml_logits_processor_core::register_module(py, m)?;
-    
+    crate::py_xml_schema_validator::register_module(py, m)?;
-    // Register the XmlLogitsProcessor class
-    crate::logits_processor::register_module(py, m)?;
-    
     Ok(())
 }

lib/xml_schema_validator/src/token/mod.rs

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

lib/xml_schema_validator/src/token/text_content.rs

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

lib/xml_schema_validator/test_requirements.txt

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

lib/xml_schema_validator/tests/conftest.py

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

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -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()

lib/xml_schema_validator/xml_schema_validator/__init__.py

@@ -1,6 +1,8 @@
 import torch
+import sys
 from transformers import LogitsProcessor
-from xml_schema_validator import XmlSchemaValidator
+from transformers import AutoTokenizer
+import xml_schema_validator
 
 class XmlLogitsProcessor(LogitsProcessor):
     """
@@ -10,7 +12,7 @@ class XmlLogitsProcessor(LogitsProcessor):
     by setting their logits to negative infinity.
     """
     
-    def __init__(self, tokenizer, schema_text: str):
+    def __init__(self, tokenizer: AutoTokenizer, schema_text: str):
         """
         Initialize the processor with a schema and tokenizer.
         
@@ -18,8 +20,13 @@ class XmlLogitsProcessor(LogitsProcessor):
             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
-        self.schema_validator = XmlSchemaValidator(schema_text)
+        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
 
@@ -48,37 +55,21 @@ class XmlLogitsProcessor(LogitsProcessor):
             generated_ids = current_ids[self.prompt_length:]
             generated_text = self.tokenizer.decode(generated_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)
+            batch_processor = self.core.copy()
-            batch_validator = self.schema_validator.copy()
 
-            print("evaluate tokens continuing:", generated_text)
-            
             if generated_text:
-                valid, msg = batch_validator.append(generated_text)
+                batch_processor.append(generated_text) 
-                if not valid:
+                if batch_processor.eof():
-                    print("current generated text invalid, only accept eos_token_id")
+                    vocab_size = scores.shape[-1]
-                    eos_token_id = self.tokenizer.eos_token_id
+                    valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
-                    if eos_token_id is not None:
+                    valid_tokens_mask[self.eos_token_id] = True
-                        valid_tokens_mask[eos_token_id] = True
+                    invalid_tokens_mask = ~valid_tokens_mask
-                        invalid_tokens_mask = ~valid_tokens_mask
+                    scores[batch_idx, invalid_tokens_mask] = float('-inf')
-                        scores[batch_idx, invalid_tokens_mask] = float('-inf')
+                    continue
-                        continue
+
-            
+            for token in batch_processor.get_invalid_tokens():
-            # Rest of the method remains the same
+                if token < scores.shape[1]:
-            for token_idx in range(vocab_size):
+                    scores[batch_idx, token] = float('-inf')
-                token_validator = batch_validator.copy()
+
-                token_text = self.tokenizer.decode([token_idx])
-                valid, msg = token_validator.append(token_text)
-                #print("token:", token_text, "valid:", valid)
-                if valid:
-                    #print("token:", generated_text + token_text)
-                    valid_tokens_mask[token_idx] = True
-            
-            invalid_tokens_mask = ~valid_tokens_mask
-            scores[batch_idx, invalid_tokens_mask] = float('-inf')
-        
         return scores

sia/llm_engine/qwq_llm_engine.py

@@ -2,12 +2,12 @@ from pathlib import Path
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
 from typing import Callable, Iterator, Optional, List
+from xml_schema_validator import XmlLogitsProcessor
 import json
 import torch
 
 from . import LlmEngine
 from .. import util
-from .xml_logits_processor import XmlLogitsProcessor
 
 class QwQLlmEngine(LlmEngine):