Update multipurpose_chatbot/engines/transformers_engine.py

multipurpose_chatbot/engines/transformers_engine.py

@@ -429,6 +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)
         with torch.no_grad():
             inputs = self.tokenizer(prompt, return_tensors='pt')
             num_tokens = inputs.input_ids.size(1)
@@ -448,12 +449,16 @@ class TransformersEngine(BaseEngine):
             for token in generator:
                 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")
                 num_tokens += 1
                 print(f"{response}", end='\r')
                 sys.stdout.flush()
                 yield response, num_tokens
             
             if response is not None:
+                if "<|im_start|>assistant\n" in response:
+                    response = response.split("<|im_start|>assistant\n")
                 full_text = prompt + response
                 num_tokens = len(self.tokenizer.encode(full_text))
                 yield response, num_tokens