Fix start of assistant message detection

lib/xml_schema_validator/python/xml_schema_validator/__init__.py

@@ -37,8 +37,25 @@ class XmlLogitsProcessor(LogitsProcessor):
         else:
             raise ValueError("Either schema_text or core must be provided")
         
-        self.prompt_length = None
+        # Find the assistant start marker in the chat template
-        self.is_first_call = True
+        # You can also manually override this by setting a specific marker if needed
+        self.assistant_start_marker = self._get_assistant_start_marker(tokenizer)
+
+    def _get_assistant_start_marker(self, tokenizer):
+        """
+        Extract the assistant start marker from the tokenizer's chat template.
+        
+        Args:
+            tokenizer: The tokenizer to extract the marker from
+            
+        Returns:
+            The marker string that indicates the start of the assistant's response
+        """
+        return tokenizer.apply_chat_template(
+            [{"role": "assistant", "content": ""},],
+            tokenize=False,
+            add_generation_prompt=False,
+        )
 
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
         """
@@ -53,19 +70,34 @@ 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 only the generated portion of the text
+            # Decode the entire sequence so far
             current_ids = input_ids[batch_idx]
-            generated_ids = current_ids[self.prompt_length:]
+            full_text = self.tokenizer.decode(current_ids)
-            generated_text = self.tokenizer.decode(generated_ids)
             
-            # Create a mask to track which tokens are valid
+            # Find the last occurrence of the assistant start marker
+            last_marker_pos = full_text.rfind(self.assistant_start_marker)
+            
+            if last_marker_pos == -1:
+                # If primary marker not found, try looking for a second common marker as fallback
+                fallback_markers = ["<|assistant|>", "\nAssistant: ", "\nA: "]
+                for fallback in fallback_markers:
+                    last_marker_pos = full_text.rfind(fallback)
+                    if last_marker_pos != -1:
+                        # Found a fallback marker
+                        self.assistant_start_marker = fallback  # Update for future calls
+                        break
+                        
+                # If still not found, we can't determine where the assistant content starts
+                if last_marker_pos == -1:
+                    continue
+                
+            # Extract the assistant's response by taking text after the marker plus its length
+            start_pos = last_marker_pos + len(self.assistant_start_marker)
+            generated_text = full_text[start_pos:]
+            
+            # Create a processor copy to track which tokens are valid
             batch_processor = self.core.copy()
 
             if generated_text:
@@ -106,6 +138,8 @@ class XmlLogitsProcessor(LogitsProcessor):
         # Create a new instance using the existing core
         cloned_core = self.core.copy()
         cloned = XmlLogitsProcessor(self.tokenizer, core=cloned_core)
-        cloned.prompt_length = self.prompt_length
+        
-        cloned.is_first_call = self.is_first_call
+        # Copy all state
+        cloned.assistant_start_marker = self.assistant_start_marker
+        
         return cloned

sia/llm_engine/qwq_llm_engine.py

@@ -27,9 +27,6 @@ class QwQLlmEngine(LlmEngine):
         """
         self._temperature = temperature
 
-        with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
-            tokenizer_config = json.load(f)
-
         quantization_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_compute_dtype=torch.bfloat16,
@@ -47,7 +44,6 @@ class QwQLlmEngine(LlmEngine):
 
         self._tokenizer = AutoTokenizer.from_pretrained(
             model_path,
-            **tokenizer_config,
         )
 
         self._pipeline = pipeline(