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