From 8fb30c8ed0cc321c07f3e183a940619b511fce3c Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Wed, 9 Apr 2025 11:59:26 +0200 Subject: [PATCH] Optimized hot-loop for logits processor --- .dockerignore | 6 +- Dockerfile | 2 +- lib/xml_schema_validator/.dockerignore | 2 - lib/xml_schema_validator/Cargo.toml | 6 +- lib/xml_schema_validator/pyproject.toml | 13 +- lib/xml_schema_validator/src/lib.rs | 8 +- .../src/logits_processor.rs | 174 ------------------ .../src/py_xml_logits_processor_core.rs | 58 ++++++ .../src/py_xml_schema_validator.rs | 42 +++++ lib/xml_schema_validator/src/python.rs | 38 +--- lib/xml_schema_validator/src/token/mod.rs | 7 + .../src/token/text_content.rs | 13 +- .../test_requirements.txt | 3 + lib/xml_schema_validator/tests/conftest.py | 9 + .../tests/test_logits_processor.py | 71 +++++++ .../xml_schema_validator/__init__.py | 57 +++--- sia/llm_engine/qwq_llm_engine.py | 2 +- 17 files changed, 239 insertions(+), 272 deletions(-) delete mode 100644 lib/xml_schema_validator/.dockerignore delete mode 100644 lib/xml_schema_validator/src/logits_processor.rs create mode 100644 lib/xml_schema_validator/src/py_xml_logits_processor_core.rs create mode 100644 lib/xml_schema_validator/src/py_xml_schema_validator.rs create mode 100644 lib/xml_schema_validator/test_requirements.txt create mode 100644 lib/xml_schema_validator/tests/conftest.py create mode 100644 lib/xml_schema_validator/tests/test_logits_processor.py rename sia/llm_engine/xml_logits_processor.py => lib/xml_schema_validator/xml_schema_validator/__init__.py (50%) diff --git a/.dockerignore b/.dockerignore index e346f21..9facfd9 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,2 +1,4 @@ -.env -./model/ \ No newline at end of file +.env +./model/ +**/target/ +**/Cargo.lock \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index 92681cd..aafccaf 100644 --- a/Dockerfile +++ b/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"] \ No newline at end of file diff --git a/lib/xml_schema_validator/.dockerignore b/lib/xml_schema_validator/.dockerignore deleted file mode 100644 index 2630b86..0000000 --- a/lib/xml_schema_validator/.dockerignore +++ /dev/null @@ -1,2 +0,0 @@ -target/ -Cargo.lock \ No newline at end of file diff --git a/lib/xml_schema_validator/Cargo.toml b/lib/xml_schema_validator/Cargo.toml index 4cf4733..6ba8f0b 100644 --- a/lib/xml_schema_validator/Cargo.toml +++ b/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" - -[build-dependencies] -xml_schema_generator = { version = "0.6.18", features = ["env_logger"] } \ No newline at end of file +thiserror = "1.0.38" \ No newline at end of file diff --git a/lib/xml_schema_validator/pyproject.toml b/lib/xml_schema_validator/pyproject.toml index 7587057..3026f43 100644 --- a/lib/xml_schema_validator/pyproject.toml +++ b/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"] diff --git a/lib/xml_schema_validator/src/lib.rs b/lib/xml_schema_validator/src/lib.rs index 7af247e..1c697ea 100644 --- a/lib/xml_schema_validator/src/lib.rs +++ b/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![]; diff --git a/lib/xml_schema_validator/src/logits_processor.rs b/lib/xml_schema_validator/src/logits_processor.rs deleted file mode 100644 index 7bf684c..0000000 --- a/lib/xml_schema_validator/src/logits_processor.rs +++ /dev/null @@ -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>, - /// Cached EOS token ID - eos_token_id: RwLock>, -} - -#[pymethods] -impl XmlLogitsProcessor { - #[new] - fn new(py: Python<'_>, tokenizer: PyObject, schema_text: &str) -> PyResult { - // 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::(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 { - // 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::() { - 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::()?; - Ok(()) -} \ No newline at end of file diff --git a/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs b/lib/xml_schema_validator/src/py_xml_logits_processor_core.rs new file mode 100644 index 0000000..35055cf --- /dev/null +++ b/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 { + 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> { + 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::>()) + } + + fn eof(&self) -> PyResult { + Ok(self.xml_schema_validator.eof()) + } + + fn copy(&self, py: Python<'_>) -> PyResult { + 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::()?; + Ok(()) +} \ No newline at end of file diff --git a/lib/xml_schema_validator/src/py_xml_schema_validator.rs b/lib/xml_schema_validator/src/py_xml_schema_validator.rs new file mode 100644 index 0000000..0849e42 --- /dev/null +++ b/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 { + 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)> { + match self.validator.append(fragment) { + Ok(()) => PyResult::Ok((true, None)), + Err(e) => PyResult::Ok((false, Some(e.to_string()))), + } + } + + fn eof(&self) -> PyResult { + Ok(self.validator.eof()) + } + + fn copy(&self) -> PyResult { + 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::()?; + Ok(()) +} \ No newline at end of file diff --git a/lib/xml_schema_validator/src/python.rs b/lib/xml_schema_validator/src/python.rs index 6b64435..d8be3ea 100644 --- a/lib/xml_schema_validator/src/python.rs +++ b/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 { - 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)> { - match self.validator.append(fragment) { - Ok(()) => PyResult::Ok((true, None)), - Err(e) => PyResult::Ok((false, Some(e.to_string()))), - } - } - - fn copy(&self) -> PyResult { - 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::()?; - - // 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(()) } \ No newline at end of file diff --git a/lib/xml_schema_validator/src/token/mod.rs b/lib/xml_schema_validator/src/token/mod.rs index 0e6dbfa..09188ed 100644 --- a/lib/xml_schema_validator/src/token/mod.rs +++ b/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, + } + } } \ No newline at end of file diff --git a/lib/xml_schema_validator/src/token/text_content.rs b/lib/xml_schema_validator/src/token/text_content.rs index 033776d..3434385 100644 --- a/lib/xml_schema_validator/src/token/text_content.rs +++ b/lib/xml_schema_validator/src/token/text_content.rs @@ -5,22 +5,12 @@ use crate::*; #[derive(Clone, Debug)] pub struct TextContent { element: Arc, - buffer: String, } impl TextContent { pub fn new(element: Arc) -> Self { Self { element, - buffer: String::new(), - } - } - - /// Create a new TextContent with the given initial buffer - pub fn with_buffer(element: Arc, 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(), ))] } } diff --git a/lib/xml_schema_validator/test_requirements.txt b/lib/xml_schema_validator/test_requirements.txt new file mode 100644 index 0000000..fcacc61 --- /dev/null +++ b/lib/xml_schema_validator/test_requirements.txt @@ -0,0 +1,3 @@ +torch>=2.0.0 +transformers>=4.30.0 +. \ No newline at end of file diff --git a/lib/xml_schema_validator/tests/conftest.py b/lib/xml_schema_validator/tests/conftest.py new file mode 100644 index 0000000..9798367 --- /dev/null +++ b/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"]) \ No newline at end of file diff --git a/lib/xml_schema_validator/tests/test_logits_processor.py b/lib/xml_schema_validator/tests/test_logits_processor.py new file mode 100644 index 0000000..327179c --- /dev/null +++ b/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 = """ + + + + + + + + +""" + + 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 message by the user. So if the user says 'hello world', you response hello world"}, + {"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() \ No newline at end of file diff --git a/sia/llm_engine/xml_logits_processor.py b/lib/xml_schema_validator/xml_schema_validator/__init__.py similarity index 50% rename from sia/llm_engine/xml_logits_processor.py rename to lib/xml_schema_validator/xml_schema_validator/__init__.py index c58ca82..1586866 100644 --- a/sia/llm_engine/xml_logits_processor.py +++ b/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_validator = self.schema_validator.copy() + batch_processor = self.core.copy() - print("evaluate tokens continuing:", generated_text) - if generated_text: - valid, msg = batch_validator.append(generated_text) - if not valid: - print("current generated text invalid, only accept eos_token_id") - eos_token_id = self.tokenizer.eos_token_id - if eos_token_id is not None: - valid_tokens_mask[eos_token_id] = True - invalid_tokens_mask = ~valid_tokens_mask - scores[batch_idx, invalid_tokens_mask] = float('-inf') - continue - - # Rest of the method remains the same - for token_idx in range(vocab_size): - 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') - + 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 \ 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 b4413f9..498d42a 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/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):