Allow stopping of inference

sia/hf_llm_engine.py

@@ -1,6 +1,6 @@
 from huggingface_hub import InferenceClient
 from transformers import AutoTokenizer, AutoConfig
-from typing import Iterator, Optional
+from typing import Iterator, Optional, Callable
 
 from .llm_engine import LlmEngine
 
@@ -30,36 +30,39 @@ class HfLlmEngine(LlmEngine):
         self._config = AutoConfig.from_pretrained(model, token=api_token)
         self._client = InferenceClient(token=api_token)
                         
-    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.
 
         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.
         """
-        token_count=self.token_count(system_prompt, main_context)
-
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
         ]
         
-        def stream_wrapper():
-            stream = self._client.chat_completion(
-                model=self._model,
-                messages=messages,
-                temperature=self._temperature,
-                stream=True
-            )
-            
+        stream = self._client.chat_completion(
+            model=self._model,
+            messages=messages,
+            temperature=self._temperature,
+            stream=True
+        )
+        
+        try:
             for response in stream:
+                if should_stop():
+                    stream.close()
+                    break
                 if content := response.choices[0].delta.content:
                     yield content
-        return stream_wrapper()
+        finally:
+            stream.close()
 
     def token_count(self, system_prompt: str, main_context: str) -> int:
         messages = [