From 5b8f04be81000f8f7d4959ec9d93ac6cd42f0d0d Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@materialise.be>
Date: Wed, 9 Apr 2025 16:43:27 +0200
Subject: [PATCH] fix copy bug

---
 .../python/xml_schema_validator/__init__.py   | 19 ++++++++--
 sia/__main__.py                               |  1 +
 sia/llm_engine/local_llm_engine.py            | 35 +++++++++++++------
 sia/llm_engine/qwq_llm_engine.py              |  2 +-
 4 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
index 42182ca..7bbc9e1 100644
--- a/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
+++ b/lib/xml_schema_validator/python/xml_schema_validator/__init__.py
@@ -26,7 +26,8 @@ class XmlLogitsProcessor(LogitsProcessor):
         items = dict()
         for token, id in vocab.items():
             items[id] = token
-        self.core = XmlLogitsProcessorCore(items, schema_text)
+        if schema_text:
+            self.core = XmlLogitsProcessorCore(items, schema_text)
         self.prompt_length = None
         self.is_first_call = True
 
@@ -72,4 +73,18 @@ class XmlLogitsProcessor(LogitsProcessor):
                 if token < scores.shape[1]:
                     scores[batch_idx, token] = float('-inf')
 
-        return scores
\ No newline at end of file
+        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
\ No newline at end of file
diff --git a/sia/__main__.py b/sia/__main__.py
index bc2b3c0..04eccd9 100644
--- a/sia/__main__.py
+++ b/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,
             )
             
diff --git a/sia/llm_engine/local_llm_engine.py b/sia/llm_engine/local_llm_engine.py
index 887a000..439f747 100644
--- a/sia/llm_engine/local_llm_engine.py
+++ b/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
diff --git a/sia/llm_engine/qwq_llm_engine.py b/sia/llm_engine/qwq_llm_engine.py
index 498d42a..9f29fe7 100644
--- a/sia/llm_engine/qwq_llm_engine.py
+++ b/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,