From 2e08be4b756f1227b7c7f538cd2aa67f6b1ae88b Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Mon, 7 Apr 2025 14:31:56 +0000 Subject: [PATCH] Validator works for start --- lib/xml_schema_validator/src/lib.rs | 5 +-- lib/xml_schema_validator/src/python.rs | 10 +++-- .../src/token/element_close_end.rs | 4 +- .../src/token/element_close_name.rs | 20 ++++----- .../src/token/element_close_start.rs | 18 ++++---- .../src/token/element_open_end.rs | 14 +++---- .../src/token/element_open_name.rs | 20 ++++----- .../src/token/element_open_start.rs | 10 ++--- .../src/token/end_of_file.rs | 2 +- lib/xml_schema_validator/src/token/mod.rs | 2 +- .../src/token/start_of_file.rs | 8 ++-- .../src/token/text_content.rs | 16 ++++---- .../src/token/whitespace.rs | 10 ++--- sia/__main__.py | 1 + sia/llm_engine/qwq_llm_engine.py | 9 ++-- sia/llm_engine/xml_logits_processor.py | 41 ++++++++++++++----- 16 files changed, 108 insertions(+), 82 deletions(-) diff --git a/lib/xml_schema_validator/src/lib.rs b/lib/xml_schema_validator/src/lib.rs index ea23073..2f400f0 100644 --- a/lib/xml_schema_validator/src/lib.rs +++ b/lib/xml_schema_validator/src/lib.rs @@ -14,7 +14,7 @@ use token::*; /// An XML validator that checks XML fragments against an XSD schema #[derive(Clone, Debug)] pub struct XmlSchemaValidator { - current_tokens: Vec>, + current_tokens: Vec, } /// The Result type used throughout this crate @@ -27,7 +27,7 @@ impl XmlSchemaValidator { pub fn new(schema_text: &str) -> Result { let schema: schema::XsSchema = quick_xml::de::from_str(schema_text)?; let current_tokens = schema.xs_element.iter().map(|element| { - Box::new(Token::StartOfFile(StartOfFile::new(Box::new(element.clone())))) + Token::StartOfFile(StartOfFile::new(element.clone())) }).collect(); Ok(Self { @@ -45,7 +45,6 @@ impl XmlSchemaValidator { /// Append a single character to the validator fn append_char(&mut self, c: char) -> Result<()> { - println!("Appending char: {}", c); let mut new_tokens = vec![]; for token in self.current_tokens.clone() { new_tokens.append(&mut token.append(c)); diff --git a/lib/xml_schema_validator/src/python.rs b/lib/xml_schema_validator/src/python.rs index f4051eb..6654226 100644 --- a/lib/xml_schema_validator/src/python.rs +++ b/lib/xml_schema_validator/src/python.rs @@ -2,7 +2,7 @@ use pyo3::prelude::*; use pyo3::exceptions::PyValueError; /// Python wrapper for XmlValidator -#[pyclass(unsendable)] +#[pyclass] struct XmlSchemaValidator { validator: crate::XmlSchemaValidator, } @@ -17,9 +17,11 @@ impl XmlSchemaValidator { Ok(Self { validator }) } - fn append(&mut self, fragment: &str) -> PyResult<()> { - self.validator.append(fragment) - .map_err(|e| PyValueError::new_err(format!("Validation error: {}", e))) + 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 { diff --git a/lib/xml_schema_validator/src/token/element_close_end.rs b/lib/xml_schema_validator/src/token/element_close_end.rs index ababfe1..74566f9 100644 --- a/lib/xml_schema_validator/src/token/element_close_end.rs +++ b/lib/xml_schema_validator/src/token/element_close_end.rs @@ -11,9 +11,9 @@ impl ElementCloseEnd { } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { if '>' == c { - vec![Box::new(Token::EndOfFile(EndOfFile::new()))] + vec![Token::EndOfFile(EndOfFile::new())] } else { vec![] } diff --git a/lib/xml_schema_validator/src/token/element_close_name.rs b/lib/xml_schema_validator/src/token/element_close_name.rs index 57802ba..8b4ac89 100644 --- a/lib/xml_schema_validator/src/token/element_close_name.rs +++ b/lib/xml_schema_validator/src/token/element_close_name.rs @@ -5,21 +5,21 @@ use crate::*; #[derive(Clone, Debug)] pub struct ElementCloseName { /// The element that was previously opened and is now being closed - element: Box, + element: schema::XsElement, /// Current buffer of characters processed for the element name buffer: String, } impl ElementCloseName { /// Create a new instance with the given buffer - pub fn new(element: Box, buffer: String) -> Self { + pub fn new(element: schema::XsElement, buffer: String) -> Self { Self { element, buffer, } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { let new_buffer = self.buffer.clone() + &c.to_string(); // If still building the element name @@ -29,19 +29,19 @@ impl ElementCloseName { // Now we're expecting '>' to close the tag if c.is_whitespace() { // Handle whitespace after the element name - vec![Box::new(Token::Whitespace(Whitespace::new( - vec![Box::new(Token::ElementCloseEnd(ElementCloseEnd::new()))] - )))] + vec![Token::Whitespace(Whitespace::new( + vec![Token::ElementCloseEnd(ElementCloseEnd::new())] + ))] } else { // Directly transition to the closed state - vec![Box::new(Token::ElementCloseEnd(ElementCloseEnd::new()))] + vec![Token::ElementCloseEnd(ElementCloseEnd::new())] } } else { // Still building the element name - vec![Box::new(Token::ElementCloseName(ElementCloseName::new( - Box::new((*self.element).clone()), + vec![Token::ElementCloseName(ElementCloseName::new( + self.element.clone(), new_buffer, - )))] + ))] } } else { // The character doesn't match what we expect for this element name diff --git a/lib/xml_schema_validator/src/token/element_close_start.rs b/lib/xml_schema_validator/src/token/element_close_start.rs index bb13a5a..328f4cc 100644 --- a/lib/xml_schema_validator/src/token/element_close_start.rs +++ b/lib/xml_schema_validator/src/token/element_close_start.rs @@ -4,30 +4,30 @@ use crate::*; #[derive(Clone, Debug)] pub struct ElementCloseStart { /// The element that was previously opened and is now being closed - element: Box, + element: schema::XsElement, has_slash: bool, } impl ElementCloseStart { - pub fn new(element: Box) -> Self { + pub fn new(element: schema::XsElement) -> Self { Self { element, has_slash: false, } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { if !self.has_slash && '/' == c { - vec![Box::new(Token::ElementCloseStart(Self { - element: Box::new((*self.element).clone()), + vec![Token::ElementCloseStart(Self { + element: self.element.clone(), has_slash: true, - }))] + })] } else { if self.element.name.starts_with(c) { - vec![Box::new(Token::ElementCloseName(ElementCloseName::new( - Box::new((*self.element).clone()), + vec![Token::ElementCloseName(ElementCloseName::new( + self.element.clone(), c.to_string(), - )))] + ))] } else { vec![] } diff --git a/lib/xml_schema_validator/src/token/element_open_end.rs b/lib/xml_schema_validator/src/token/element_open_end.rs index 680101e..5f7387e 100644 --- a/lib/xml_schema_validator/src/token/element_open_end.rs +++ b/lib/xml_schema_validator/src/token/element_open_end.rs @@ -3,17 +3,17 @@ use crate::*; /// Represents a fully opened XML element, containing the schema element reference #[derive(Clone, Debug)] pub struct ElementOpenEnd { - element: Box, + element: schema::XsElement, } impl ElementOpenEnd { - pub fn new(element: Box) -> Self { + pub fn new(element: schema::XsElement) -> Self { Self { element, } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { // Check if the element allows text content (mixed="true") let allows_text = self.element.xs_complex_type.mixed .as_ref() @@ -22,11 +22,11 @@ impl ElementOpenEnd { if allows_text { if c.is_whitespace() { - vec![Box::new(Token::Whitespace(Whitespace::new( - vec![Box::new(Token::TextContent(TextContent::new(Box::new((*self.element).clone()))))] - )))] + vec![Token::Whitespace(Whitespace::new( + vec![Token::TextContent(TextContent::new(self.element.clone()))] + ))] } else { - vec![Box::new(Token::TextContent(TextContent::new(Box::new((*self.element).clone()))))] + vec![Token::TextContent(TextContent::new(self.element.clone()))] } } else { vec![] diff --git a/lib/xml_schema_validator/src/token/element_open_name.rs b/lib/xml_schema_validator/src/token/element_open_name.rs index d2194c5..25bd28e 100644 --- a/lib/xml_schema_validator/src/token/element_open_name.rs +++ b/lib/xml_schema_validator/src/token/element_open_name.rs @@ -4,11 +4,11 @@ use crate::*; #[derive(Clone, Debug)] pub struct ElementOpenName { buffer: String, - element: Box, + element: schema::XsElement, } impl ElementOpenName { - pub fn new(element: Box, buffer: String) -> Result { + pub fn new(element: schema::XsElement, buffer: String) -> Result { if element.name.starts_with(buffer.as_str()) { Ok(Self { buffer, @@ -19,19 +19,19 @@ impl ElementOpenName { } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { let buffer = self.buffer.clone() + &c.to_string(); if self.element.name.starts_with(buffer.as_str()) { - vec![Box::new(Token::ElementName(ElementOpenName { + vec![Token::ElementName(ElementOpenName { buffer, - element: Box::new((*self.element).clone()), - }))] + element: self.element.clone(), + })] } else if c.is_whitespace() { - vec![Box::new(Token::Whitespace(Whitespace::new( - vec![Box::new(Token::ElementOpen(ElementOpenEnd::new(Box::new((*self.element).clone()))))] - )))] + vec![Token::Whitespace(Whitespace::new( + vec![Token::ElementOpen(ElementOpenEnd::new(self.element.clone()))] + ))] } else if '>' == c { - vec![Box::new(Token::ElementOpen(ElementOpenEnd::new(Box::new((*self.element).clone()))))] + vec![Token::ElementOpen(ElementOpenEnd::new(self.element.clone()))] } else { vec![] } diff --git a/lib/xml_schema_validator/src/token/element_open_start.rs b/lib/xml_schema_validator/src/token/element_open_start.rs index 612bf0d..b832076 100644 --- a/lib/xml_schema_validator/src/token/element_open_start.rs +++ b/lib/xml_schema_validator/src/token/element_open_start.rs @@ -3,20 +3,20 @@ use crate::*; /// Represents an XML element that is in the process of being opened, containing the text buffer and schema element reference #[derive(Clone, Debug)] pub struct ElementOpenStart { - element: Box, + element: schema::XsElement, } impl ElementOpenStart { - pub fn new(element: Box) -> Self { + pub fn new(element: schema::XsElement) -> Self { Self { element, } } - pub fn append(&self, c: char) -> Vec> { - let element_name = ElementOpenName::new(Box::new((*self.element).clone()), c.to_string()); + pub fn append(&self, c: char) -> Vec { + let element_name = ElementOpenName::new(self.element.clone(), c.to_string()); if let Ok(element_name) = element_name { - vec![Box::new(Token::ElementName(element_name))] + vec![Token::ElementName(element_name)] } else { vec![] } diff --git a/lib/xml_schema_validator/src/token/end_of_file.rs b/lib/xml_schema_validator/src/token/end_of_file.rs index 3f50925..4b17cfb 100644 --- a/lib/xml_schema_validator/src/token/end_of_file.rs +++ b/lib/xml_schema_validator/src/token/end_of_file.rs @@ -10,7 +10,7 @@ impl EndOfFile { Self {} } - pub fn append(&self, _c: char) -> Vec> { + pub fn append(&self, _c: char) -> Vec { vec![] } } \ 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 95be565..9f8def0 100644 --- a/lib/xml_schema_validator/src/token/mod.rs +++ b/lib/xml_schema_validator/src/token/mod.rs @@ -36,7 +36,7 @@ pub enum Token { } impl Token { - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { match self { Token::ElementCloseEnd(element_close_end) => element_close_end.append(c), Token::ElementCloseName(element_close_name) => element_close_name.append(c), diff --git a/lib/xml_schema_validator/src/token/start_of_file.rs b/lib/xml_schema_validator/src/token/start_of_file.rs index ae7176f..8e4bc64 100644 --- a/lib/xml_schema_validator/src/token/start_of_file.rs +++ b/lib/xml_schema_validator/src/token/start_of_file.rs @@ -3,17 +3,17 @@ use crate::*; /// Represents the initial state at the start of an XML file, containing a reference to the root element schema #[derive(Clone, Debug)] pub struct StartOfFile { - element: Box, + element: schema::XsElement, } impl StartOfFile { - pub fn new(element: Box) -> Self { + pub fn new(element: schema::XsElement) -> Self { Self { element } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { if '<' == c { - vec![Box::new(Token::ElementStart(ElementOpenStart::new(Box::new((*self.element).clone()))))] + vec![Token::ElementStart(ElementOpenStart::new(self.element.clone()))] } else { vec![] } diff --git a/lib/xml_schema_validator/src/token/text_content.rs b/lib/xml_schema_validator/src/token/text_content.rs index 1ab25dc..1b80991 100644 --- a/lib/xml_schema_validator/src/token/text_content.rs +++ b/lib/xml_schema_validator/src/token/text_content.rs @@ -3,12 +3,12 @@ use crate::*; /// Represents the text content within an XML element #[derive(Clone, Debug)] pub struct TextContent { - element: Box, + element: schema::XsElement, buffer: String, } impl TextContent { - pub fn new(element: Box) -> Self { + pub fn new(element: schema::XsElement) -> Self { Self { element, buffer: String::new(), @@ -16,21 +16,21 @@ impl TextContent { } /// Create a new TextContent with the given initial buffer - pub fn with_buffer(element: Box, buffer: String) -> Self { + pub fn with_buffer(element: schema::XsElement, buffer: String) -> Self { Self { element, buffer, } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { if c == '<' { - vec![Box::new(Token::ElementCloseStart(ElementCloseStart::new(Box::new((*self.element).clone()))))] + vec![Token::ElementCloseStart(ElementCloseStart::new(self.element.clone()))] } else { - vec![Box::new(Token::TextContent(TextContent::with_buffer( - Box::new((*self.element).clone()), + vec![Token::TextContent(TextContent::with_buffer( + self.element.clone(), self.buffer.clone() + &c.to_string(), - )))] + ))] } } } \ No newline at end of file diff --git a/lib/xml_schema_validator/src/token/whitespace.rs b/lib/xml_schema_validator/src/token/whitespace.rs index b4964fb..5ea4376 100644 --- a/lib/xml_schema_validator/src/token/whitespace.rs +++ b/lib/xml_schema_validator/src/token/whitespace.rs @@ -3,21 +3,21 @@ use crate::*; /// Represents one or more whitespace characters #[derive(Clone, Debug)] pub struct Whitespace { - next: Vec> + next: Vec } impl Whitespace { - pub fn new(next: Vec>) -> Self { + pub fn new(next: Vec) -> Self { Self { next } } - pub fn append(&self, c: char) -> Vec> { + pub fn append(&self, c: char) -> Vec { if c.is_whitespace() { - vec![Box::new(Token::Whitespace(Self { + vec![Token::Whitespace(Self { next: self.next.clone() - }))] + })] } else { let mut tokens = vec![]; for token in self.next.iter() { diff --git a/sia/__main__.py b/sia/__main__.py index eb6491e..bc2b3c0 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -67,6 +67,7 @@ class Main: config.qwq_model, config.qwq_temperature, config.qwq_token_limit, + self._action_schema, ) if not self._llms: diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py index 94740be..b4413f9 100644 --- a/sia/llm_engine/qwq_llm_engine.py +++ b/sia/llm_engine/qwq_llm_engine.py @@ -29,8 +29,6 @@ class QwQLlmEngine(LlmEngine): """ 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) @@ -62,6 +60,11 @@ class QwQLlmEngine(LlmEngine): return_full_text=False, ) + if xml_schema_text: + self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text) + else: + self._logits_processor = None + def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]: """ @@ -100,7 +103,7 @@ class QwQLlmEngine(LlmEngine): } if self._logits_processor: - generation_kwargs["logits_processor"] = self.logits_processor + generation_kwargs["logits_processor"] = [self._logits_processor] generation_thread = Thread( target=self._pipline, diff --git a/sia/llm_engine/xml_logits_processor.py b/sia/llm_engine/xml_logits_processor.py index b7b8553..c58ca82 100644 --- a/sia/llm_engine/xml_logits_processor.py +++ b/sia/llm_engine/xml_logits_processor.py @@ -20,7 +20,9 @@ class XmlLogitsProcessor(LogitsProcessor): """ self.tokenizer = tokenizer self.schema_validator = XmlSchemaValidator(schema_text) - + self.prompt_length = None + self.is_first_call = True + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: """ Process logits to mask invalid XML tokens. @@ -34,29 +36,48 @@ class XmlLogitsProcessor(LogitsProcessor): """ batch_size, _ = input_ids.shape + # If this is the first call, store the prompt length + if self.is_first_call: + self.prompt_length = input_ids.shape[1] + self.is_first_call = False + # For each sequence in the batch for batch_idx in range(batch_size): - # Get the current text generated so far + # Get only the generated portion of the text current_ids = input_ids[batch_idx] - current_text = self.tokenizer.decode(current_ids) + 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() + + print("evaluate tokens continuing:", generated_text) - # For each possible next token + 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): - # 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_validator = batch_validator.copy() token_text = self.tokenizer.decode([token_idx]) - if validator_copy.append(token_text): + 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 - # 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')