Start work on docker module

sia/llm_engine.py

@@ -1,8 +1,9 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
+from threading import Thread
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextIteratorStreamer
 import torch
 
 from . import util
-from .inference_result import InferenceResult
 
 class LlmEngine:
     def __init__(self, model_path: str):
@@ -34,33 +35,26 @@ class LlmEngine:
             self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         if model.config.pad_token_id is None:
             model.config.pad_token_id = model.config.eos_token_id
-        streamer = TextStreamer(
-            self.tokenizer,
-            skip_prompt=True
-        )
         self.pipeline = pipeline(
             "text-generation",
             model=model,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,
             device_map="auto",
-            streamer=streamer,
             return_full_text=False,
         )
 
-    def infer(self, system_prompt: str, main_context: str, action_schema: str) -> InferenceResult:
+    def infer(self, system_prompt: str, main_context: str) -> TextIteratorStreamer:
         """
         Run inference using the system prompt and main context, while validating actions against the provided XML schema.
 
         Args:
             system_prompt: The system prompt string
             main_context: The main context string after templating
-            action_schema: XML schema to validate the generated actions
 
         Returns:
-            InferenceResult: Tuple containing reasoning and actions that validate against the schema
+            TextIteratorStreamer: An iterator that yields the generated text.
         """
-        valid_elements = util.get_valid_root_elements(action_schema)
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": main_context}
@@ -68,11 +62,19 @@ class LlmEngine:
         prompt = self.tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True
         )
-        outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True)
-        generated_text = outputs[0]["generated_text"]
-        #response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip()
-        result = util.split_response(generated_text, valid_elements)
-        return result
+        streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True
+        )
+        pipeline_kwargs = dict(
+            text_inputs=prompt,
+            do_sample=True,
+            max_new_tokens=1024,
+            streamer=streamer
+        )
+        thread = Thread(target=self.pipeline, kwargs=pipeline_kwargs)
+        thread.start()
+        return util.stop_before_value(streamer, '<|eot_id|>')
 
     def finetune(self, dataset_paths: list, output_dir: str):
         """