Validator works for start

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