Optimized hot-loop for logits processor

sia/llm_engine/qwq_llm_engine.py

@@ -2,12 +2,12 @@ from pathlib import Path
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer, BitsAndBytesConfig
 from typing import Callable, Iterator, Optional, List
+from xml_schema_validator import XmlLogitsProcessor
 import json
 import torch
 
 from . import LlmEngine
 from .. import util
-from .xml_logits_processor import XmlLogitsProcessor
 
 class QwQLlmEngine(LlmEngine):
     

sia/llm_engine/xml_logits_processor.py

@@ -1,84 +0,0 @@
-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)
-        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.
-        
-        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
-        
-        # 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 only the generated portion of the text
-            current_ids = input_ids[batch_idx]
-            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)
-            
-            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):
-                token_validator = batch_validator.copy()
-                token_text = self.tokenizer.decode([token_idx])
-                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
-            
-            invalid_tokens_mask = ~valid_tokens_mask
-            scores[batch_idx, invalid_tokens_mask] = float('-inf')
-        
-        return scores