Added text streaming

sia/llm_engine.py

@@ -1,4 +1,4 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
 import torch
 
 from . import util
@@ -34,12 +34,18 @@ 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:
@@ -64,8 +70,8 @@ class LlmEngine:
         )
         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(response, valid_elements)
+        #response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip()
+        result = util.split_response(generated_text, valid_elements)
         return result
 
     def finetune(self, dataset_paths: list, output_dir: str):