Update handler.py

handler.py

@@ -1,5 +1,6 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import re
+import time
 import torch
 
 template = """Alice Gate's Persona: Alice Gate is a young, computer engineer-nerd with a knack for problem solving and a passion for technology.
@@ -26,15 +27,19 @@ class EndpointHandler():
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(
             path,
-            low_cpu_mem_usage = True, 
+            low_cpu_mem_usage = True,
             trust_remote_code = False,
-            torch_dtype = torch.float16,
+            torch_dtype = torch.float16
         ).to('cuda')
     
     def response(self, result, user_name):
         result = result.rsplit("Alice Gate:", 1)[1].split(f"{user_name}:",1)[0].strip()
-        result = re.sub('\*.*?\*', '', result)
+        parsed_result = re.sub('\*.*?\*', '', result).strip()
+        result = parsed_result if len(parsed_result) != 0 else result.replace("*","")
         result = " ".join(result.split())
+        try:
+            result = result[:[m.start() for m in re.finditer(r'[.!?]', result)][-1]+1]
+        except Exception: pass
         return {
             "message": result
         }
@@ -43,11 +48,12 @@ class EndpointHandler():
         inputs = data.pop("inputs", data)
         user_name = inputs["user_name"]
         user_input = "\n".join(inputs["user_input"])
+        prompt = template.format(
+            user_name = user_name,
+            user_input = user_input
+        )
         input_ids = self.tokenizer(
-            template.format(
-                user_name = user_name,
-                user_input = user_input
-            ),
+            prompt,
             return_tensors = "pt"
         ).to("cuda")
         generator = self.model.generate(