Update multipurpose_chatbot/engines/transformers_engine.py

multipurpose_chatbot/engines/transformers_engine.py

@@ -429,7 +429,7 @@ class TransformersEngine(BaseEngine):
         
         # ! MUST PUT INSIDE torch.no_grad() otherwise it will overflow OOM
         import sys
-        # self._model._sample = types.MethodType(NewGenerationMixin.sample_stream, self._model)
+        self._model._sample = types.MethodType(NewGenerationMixin.sample_stream, self._model)
         with torch.no_grad():
             inputs = self.tokenizer(prompt, return_tensors='pt')
             num_tokens = inputs.input_ids.size(1)
@@ -450,7 +450,7 @@ class TransformersEngine(BaseEngine):
                 out_tokens.extend(token.tolist())
                 response = self.tokenizer.decode(out_tokens)
                 if "<|im_start|>assistant\n" in response:
-                    response = response.split("<|im_start|>assistant\n")
+                    response = response.split("<|im_start|>assistant\n")[-1]
                 num_tokens += 1
                 print(f"{response}", end='\r')
                 sys.stdout.flush()
@@ -458,7 +458,7 @@ class TransformersEngine(BaseEngine):
             
             if response is not None:
                 if "<|im_start|>assistant\n" in response:
-                    response = response.split("<|im_start|>assistant\n")
+                    response = response.split("<|im_start|>assistant\n")[-1]
                 full_text = prompt + response
                 num_tokens = len(self.tokenizer.encode(full_text))
                 yield response, num_tokens