stop on stop token

app.py

@@ -50,7 +50,8 @@ def generate(
         pad_token_id=current_tokenizer.eos_token_id
     )
 
-    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)
+    streamer.stop_signal = current_tokenizer.decode(current_tokenizer.eos_token_id)
     generate_kwargs["streamer"] = streamer
 
     thread = Thread(target=current_model.generate, kwargs=generate_kwargs)
@@ -61,6 +62,9 @@ def generate(
     for new_text in streamer:
         if isinstance(new_text, torch.Tensor):
             new_text = current_tokenizer.decode(new_text)
+        if streamer.stop_signal in new_text:
+            output += new_text.split(streamer.stop_signal)[0]
+            break
         output += new_text
         yield output