Stop generation at at `\nHuman` and `\n----` in chat mode

app.py

@@ -89,6 +89,8 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
         do_sample=True,
         seed=42,
     )
+    if chat_mode:
+        generate_kwargs.update({"stop_sequences": ["\nHuman", "\n-----"]})
 
     if chat_mode and FIM_INDICATOR in prompt:
         raise ValueError("Chat mode and FIM are mutually exclusive. Choose one or the other.")
@@ -114,11 +116,15 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
     else:
         output = prompt
 
+    previous_token = ""
     for response in stream:
         if fim_mode and response.token.text =="<|endoftext|>":
             output += (suffix + "\n" + response.token.text)
+        elif chat_mode and response.token.text in ["Human", "-----"] and previous_token=="\n":
+            return output
         else:
             output += response.token.text
+        previous_token = response.token.text
         yield output
     return output