Validator works for start

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
-                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')