stop on assistant_name

app.py

@@ -61,7 +61,7 @@ def predict(
             watermark=watermark,
             temperature=temperature,
             max_new_tokens=500,
-            stop_sequences=[user_name.rstrip()],
+            stop_sequences=[user_name.rstrip(), assistant_name.rstrip()],
     )):
         if response.token.special:
             continue
@@ -69,6 +69,8 @@ def predict(
         partial_words = partial_words + response.token.text
         if partial_words.endswith(user_name.rstrip()):
             partial_words = partial_words.rstrip(user_name.rstrip())
+        if partial_words.endswith(assistant_name.rstrip()):
+            partial_words = partial_words.rstrip(assistant_name.rstrip())
 
         if i == 0:
             history.append(" " + partial_words)