Integrate logits processor with qwq (untested)

sia/llm_engine/xml_logits_processor.py

@@ -0,0 +1,63 @@
+import torch
+from transformers import LogitsProcessor
+from xml_schema_validator import XmlSchemaValidator
+
+class XmlLogitsProcessor(LogitsProcessor):
+    """
+    A LogitsProcessor that enforces valid XML according to a schema.
+    
+    This processor masks tokens that would lead to invalid XML
+    by setting their logits to negative infinity.
+    """
+    
+    def __init__(self, tokenizer, schema_text: str):
+        """
+        Initialize the processor with a schema and tokenizer.
+        
+        Args:
+            tokenizer: The tokenizer to use for decoding tokens
+            schema_text: The XSD schema text to validate against
+        """
+        self.tokenizer = tokenizer
+        self.schema_validator = XmlSchemaValidator(schema_text)
+    
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        """
+        Process logits to mask invalid XML tokens.
+        
+        Args:
+            input_ids: Current input token IDs
+            scores: Current scores (logits) for next token prediction
+            
+        Returns:
+            Processed scores with invalid tokens masked
+        """
+        batch_size, _ = input_ids.shape
+        
+        # For each sequence in the batch
+        for batch_idx in range(batch_size):
+            # Get the current text generated so far
+            current_ids = input_ids[batch_idx]
+            current_text = self.tokenizer.decode(current_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)
+            
+            # For each possible next token
+            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_text = self.tokenizer.decode([token_idx])
+                if validator_copy.append(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')
+        
+        return scores