Stop generation at \nHuman and \n---- in Chat mode (#15)

- Stop generation at at `\nHuman` and `\n----` in chat mode (76bb25ab1928288a5620db6a484d51ff7966172d)


Co-authored-by: Loubna Ben Allal <[email protected]>

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