Integrate logits processor with qwq (untested)

sia/llm_engine/qwq_llm_engine.py

@@ -1,12 +1,13 @@
 from pathlib import Path
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
-from typing import Callable, Iterator
+from typing import Callable, Iterator, Optional, List
 import json
 import torch
 
 from . import LlmEngine
 from .. import util
+from .xml_logits_processor import XmlLogitsProcessor
 
 class QwQLlmEngine(LlmEngine):
     
@@ -15,6 +16,7 @@ class QwQLlmEngine(LlmEngine):
         model_path: Path,
         temperature: float,
         token_limit: int = None,
+        xml_schema_text: Optional[str] = None,
     ):
         """
         Initialize the QwQ LLM Engine.
@@ -23,9 +25,12 @@ class QwQLlmEngine(LlmEngine):
             model_path: Local path to the model
             temperature: Sampling temperature
             token_limit: Maximum tokens to generate
+            xml_schema_text: Optional XML schema to validate against
         """
         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)
@@ -86,17 +91,22 @@ class QwQLlmEngine(LlmEngine):
             skip_prompt=True,
         )
 
+        generation_kwargs = {
+            "text_inputs": text,
+            "do_sample": True,
+            "temperature": self._temperature,
+            "max_new_tokens": self._token_limit,
+            "streamer": streamer,
+        }
+        
+        if self._logits_processor:
+            generation_kwargs["logits_processor"] = self.logits_processor
+
         generation_thread = Thread(
             target=self._pipline,
-            kwargs=dict(
-                text_inputs=text,
-                do_sample=True,
-                temperature=self._temperature,
-                max_new_tokens=self._token_limit,
-                streamer=streamer,
-            )
+            kwargs=generation_kwargs
         )
-            
+
         generation_thread.start()
 
         for text in util.stop_before_value(streamer, self._tokenizer.eos_token):

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