fix copy bug

lib/xml_schema_validator/python/xml_schema_validator/__init__.py

@@ -26,6 +26,7 @@ class XmlLogitsProcessor(LogitsProcessor):
         items = dict()
         for token, id in vocab.items():
             items[id] = token
+        if schema_text:
             self.core = XmlLogitsProcessorCore(items, schema_text)
         self.prompt_length = None
         self.is_first_call = True
@@ -73,3 +74,17 @@ class XmlLogitsProcessor(LogitsProcessor):
                     scores[batch_idx, token] = float('-inf')
 
         return scores
+
+    def copy(self):
+        """
+        Create a copy of the processor.
+
+        Returns:
+            A new instance of the processor
+        """
+        cloned = XmlLogitsProcessor(self.tokenizer, None)
+        cloned.eos_token_id = self.eos_token_id
+        cloned.core = self.core.copy()
+        cloned.prompt_length = self.prompt_length
+        cloned.is_first_call = self.is_first_call
+        return cloned

sia/__main__.py

@@ -36,6 +36,7 @@ class Main:
                 config.local_model,
                 config.local_temperature,
                 config.local_token_limit,
+                self._action_schema,
                 config.local_api_key,
             )
             

sia/llm_engine/local_llm_engine.py

@@ -1,7 +1,8 @@
 from threading import Thread
-from typing import Iterator, Optional, Callable
-
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
+from typing import Iterator, Optional, Callable
+from xml_schema_validator import XmlLogitsProcessor
+import sys
 import torch
 
 from . import LlmEngine
@@ -13,7 +14,8 @@ class LocalLlmEngine(LlmEngine):
         model_path: str,
         temperature: float,
         token_limit: int,
-        api_token: Optional[str],
+        xml_schema_text: Optional[str] = None,
+        api_token: Optional[str] = None,
     ):
         """
         Initialize the LLM Engine with a model path.
@@ -22,6 +24,7 @@ class LocalLlmEngine(LlmEngine):
             model_path: Path to the model weights to be used.
             temperature: Temperature for sampling
             token_limit: Maximum number of tokens to generate
+            xml_schema_text: Optional XML schema to validate against
             api_token: Huggingface API key
         """
         self._temperature = temperature
@@ -48,6 +51,10 @@ class LocalLlmEngine(LlmEngine):
             device_map="auto",
             return_full_text=False,
         )
+        if xml_schema_text:
+            self._logits_processor = XmlLogitsProcessor(self._tokenizer, xml_schema_text)
+        else:
+            self._logits_processor = None
 
     def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
@@ -72,16 +79,22 @@ class LocalLlmEngine(LlmEngine):
             self._tokenizer,
             skip_prompt=True
         )
-        generation_thread = Thread(target=self._pipeline, kwargs=dict(
-            text_inputs=prompt,
-            do_sample=True,
-            temperature=self._temperature,
-            max_new_tokens=self.token_limit(),
-            streamer=streamer
-        ))
+        generation_kwargs = {
+            "text_inputs": prompt,
+            "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.copy()]
+        generation_thread = Thread(
+            target=self._pipeline,
+            kwargs=generation_kwargs
+        )
         generation_thread.start()
 
-        for text in util.stop_before_value(streamer, '<|eot_id|>'):
+        for text in util.stop_before_value(streamer, self._tokenizer.eos_token):
             yield text
             if should_stop():
                 break

sia/llm_engine/qwq_llm_engine.py

@@ -103,7 +103,7 @@ class QwQLlmEngine(LlmEngine):
         }
         
         if self._logits_processor:
-            generation_kwargs["logits_processor"] = [self._logits_processor]
+            generation_kwargs["logits_processor"] = [self._logits_processor.copy()]
 
         generation_thread = Thread(
             target=self._pipline,