From 4629e751dca62509e86ca8271fa81e2ef1f3d388 Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@materialise.be>
Date: Tue, 22 Oct 2024 15:29:45 +0200
Subject: [PATCH] Added text streaming

---
 Dockerfile        |  2 +-
 sia/llm_engine.py | 12 +++++++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 513c3bf..ab047eb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -7,7 +7,7 @@ RUN pip3 install -r requirements.txt
 FROM requirements AS test
 COPY ./test/ /root/test/
 RUN mkdir -p /root/model
-CMD ["python3", "-m", "unittest", "discover", "-p", "*test.py", "-v"]
+CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"]
 
 FROM requirements
 CMD ["python3", "-m", "sia"]
\ No newline at end of file
diff --git a/sia/llm_engine.py b/sia/llm_engine.py
index ad7d87e..30c727d 100644
--- a/sia/llm_engine.py
+++ b/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):