Upload handler.py

handler.py

@@ -1,7 +1,6 @@
-from typing import Dict, Any
 import logging
-
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Dict, Any
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from peft import PeftConfig, PeftModel
 import torch.cuda
 
@@ -105,5 +104,30 @@ def generate(
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
     input_ids = input_ids.to(model.device)
 
-    # Generate the response
-    return tokenizer.decode(model.generate(input_ids, max_new_tokens=max_new_tokens))
+    # Create a TextIteratorStreamer instance
+    streamer = TextIteratorStreamer(
+        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=False
+    )
+
+    # Generate the response using TextIteratorStreamer
+    generate_kwargs = dict(
+        {"input_ids": input_ids},
+        streamer=streamer,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        top_p=top_p,
+        top_k=top_k,
+        temperature=temperature,
+        num_beams=1,
+        repetition_penalty=repetition_penalty,
+    )
+    model.generate(**generate_kwargs)
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        if "[/INST]" in "".join(outputs):
+          return "".join(outputs).replace("[/INST]","")
+        if "[INST]" in "".join(outputs):
+          return "".join(outputs).replace("[INST]","")
+    return "".join(outputs)