Allow stopping of inference

sia/local_llm_engine.py

@@ -1,5 +1,5 @@
 from threading import Thread
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Callable
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
@@ -49,13 +49,14 @@ class LocalLlmEngine(LlmEngine):
             return_full_text=False,
         )
 
-    def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
+    def infer(self, system_prompt: str, main_context: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
         """
-        Run inference using the system prompt and main context, while validating actions against the provided XML schema.
+        Run inference using the system prompt and main context.
 
         Args:
             system_prompt: The system prompt string
             main_context: The main context string after templating
+            should_stop: Callback that returns True when inference should stop
 
         Returns:
             Iterator[str]: An iterator that yields the generated text.
@@ -71,16 +72,21 @@ class LocalLlmEngine(LlmEngine):
             self._tokenizer,
             skip_prompt=True
         )
-        pipeline_kwargs = dict(
+        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
-        )
-        thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
-        thread.start()
-        return util.stop_before_value(streamer, '<|eot_id|>')
+        ))
+        generation_thread.start()
+
+        for text in util.stop_before_value(streamer, '<|eot_id|>'):
+            yield text
+            if should_stop():
+                break
+
+        generation_thread.join()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         """
@@ -112,4 +118,4 @@ class LocalLlmEngine(LlmEngine):
         if self._token_limit is not None:
             return self._token_limit
         else:
-            return self._pipeline.model.config.max_position_embeddings
+            return self._pipeline.model.config.max_position_embeddings