Validator works for start

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<Box<Token>>,
+    current_tokens: Vec<Token>,
 }
 
 /// The Result type used throughout this crate
@@ -27,7 +27,7 @@ impl XmlSchemaValidator {
     pub fn new(schema_text: &str) -> Result<Self> {
         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));

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<()> {
+    fn append(&mut self, fragment: &str) -> PyResult<(bool, Option<String>)> {
-        self.validator.append(fragment)
+        match self.validator.append(fragment) {
-            .map_err(|e| PyValueError::new_err(format!("Validation error: {}", e)))
+            Ok(()) => PyResult::Ok((true, None)),
+            Err(e) => PyResult::Ok((false, Some(e.to_string()))),
+        }
     }
 
     fn copy(&self) -> PyResult<Self> {

lib/xml_schema_validator/src/token/element_close_end.rs

@@ -11,9 +11,9 @@ impl ElementCloseEnd {
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         if '>' == c {
-            vec![Box::new(Token::EndOfFile(EndOfFile::new()))]
+            vec![Token::EndOfFile(EndOfFile::new())]
         } else {
             vec![]
         }

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<schema::XsElement>,
+    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<schema::XsElement>, buffer: String) -> Self {
+    pub fn new(element: schema::XsElement, buffer: String) -> Self {
         Self {
             element,
             buffer,
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         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![Token::Whitespace(Whitespace::new(
-                        vec![Box::new(Token::ElementCloseEnd(ElementCloseEnd::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(
+                vec![Token::ElementCloseName(ElementCloseName::new(
-                    Box::new((*self.element).clone()),
+                    self.element.clone(),
                     new_buffer,
-                )))]
+                ))]
             }
         } else {
             // The character doesn't match what we expect for this element name

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<schema::XsElement>,
+    element: schema::XsElement,
     has_slash: bool,
 }
 
 impl ElementCloseStart {
-    pub fn new(element: Box<schema::XsElement>) -> Self {
+    pub fn new(element: schema::XsElement) -> Self {
         Self {
             element,
             has_slash: false,
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         if !self.has_slash && '/' == c {
-            vec![Box::new(Token::ElementCloseStart(Self {
+            vec![Token::ElementCloseStart(Self {
-                element: Box::new((*self.element).clone()),
+                element: self.element.clone(),
                 has_slash: true,
-            }))]
+            })]
         } else {
             if self.element.name.starts_with(c) {
-                vec![Box::new(Token::ElementCloseName(ElementCloseName::new(
+                vec![Token::ElementCloseName(ElementCloseName::new(
-                    Box::new((*self.element).clone()),
+                    self.element.clone(),
                     c.to_string(),
-                )))]
+                ))]
             } else {
                 vec![]
             }

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<schema::XsElement>,
+    element: schema::XsElement,
 }
 
 impl ElementOpenEnd {
-    pub fn new(element: Box<schema::XsElement>) -> Self {
+    pub fn new(element: schema::XsElement) -> Self {
         Self {
             element,
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         // 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![Token::Whitespace(Whitespace::new(
-                    vec![Box::new(Token::TextContent(TextContent::new(Box::new((*self.element).clone()))))]
+                    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![]

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<schema::XsElement>,
+    element: schema::XsElement,
 }
 
 impl ElementOpenName {
-    pub fn new(element: Box<schema::XsElement>, buffer: String) -> Result<Self> {
+    pub fn new(element: schema::XsElement, buffer: String) -> Result<Self> {
         if element.name.starts_with(buffer.as_str()) {
             Ok(Self {
                 buffer,
@@ -19,19 +19,19 @@ impl ElementOpenName {
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         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![Token::Whitespace(Whitespace::new(
-                vec![Box::new(Token::ElementOpen(ElementOpenEnd::new(Box::new((*self.element).clone()))))]
+                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![]
         }

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<schema::XsElement>,
+    element: schema::XsElement,
 }
 
 impl ElementOpenStart {
-    pub fn new(element: Box<schema::XsElement>) -> Self {
+    pub fn new(element: schema::XsElement) -> Self {
         Self {
             element,
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
-        let element_name = ElementOpenName::new(Box::new((*self.element).clone()), c.to_string());
+        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![]
         }

lib/xml_schema_validator/src/token/end_of_file.rs

@@ -10,7 +10,7 @@ impl EndOfFile {
         Self {}
     }
 
-    pub fn append(&self, _c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, _c: char) -> Vec<Token> {
         vec![]
     }
 }

lib/xml_schema_validator/src/token/mod.rs

@@ -36,7 +36,7 @@ pub enum Token {
 }
 
 impl Token {
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         match self {
             Token::ElementCloseEnd(element_close_end) => element_close_end.append(c),
             Token::ElementCloseName(element_close_name) => element_close_name.append(c),

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<schema::XsElement>,
+    element: schema::XsElement,
 }
 
 impl StartOfFile {
-    pub fn new(element: Box<schema::XsElement>) -> Self {
+    pub fn new(element: schema::XsElement) -> Self {
         Self { element }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         if '<' == c {
-            vec![Box::new(Token::ElementStart(ElementOpenStart::new(Box::new((*self.element).clone()))))]
+            vec![Token::ElementStart(ElementOpenStart::new(self.element.clone()))]
         } else {
             vec![]
         }

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<schema::XsElement>,
+    element: schema::XsElement,
     buffer: String,
 }
 
 impl TextContent {
-    pub fn new(element: Box<schema::XsElement>) -> 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<schema::XsElement>, buffer: String) -> Self {
+    pub fn with_buffer(element: schema::XsElement, buffer: String) -> Self {
         Self {
             element,
             buffer,
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         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(
+            vec![Token::TextContent(TextContent::with_buffer(
-                Box::new((*self.element).clone()),
+                self.element.clone(),
                 self.buffer.clone() + &c.to_string(),
-            )))]
+            ))]
         }
     }
 }

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<Box<Token>>
+    next: Vec<Token>
 }
 
 impl Whitespace {
-    pub fn new(next: Vec<Box<Token>>) -> Self {
+    pub fn new(next: Vec<Token>) -> Self {
         Self {
             next
         }
     }
 
-    pub fn append(&self, c: char) -> Vec<Box<Token>> {
+    pub fn append(&self, c: char) -> Vec<Token> {
         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() {

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:

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,

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
+                token_validator = batch_validator.copy()
-                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, 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')