Update app.py

app.py

@@ -28,7 +28,7 @@ def respond(
     temperature,
     top_p,
 ):
-    # stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
+    stop_tokens = ["</s>", "[INST]", "[INST] ", "<s>", "[/INST]", "[/INST] "]
     chat_template = '<s>[INST] ' + system_message
     # for human, assistant in history:
     #     chat_template += human + ' [/INST] ' + assistant + '</s>[INST]'
@@ -54,11 +54,17 @@ def respond(
         completion_to_prompt=completion_to_prompt,
         verbose=True,
     )
-    response = ""
+    # response = ""
+    # for chunk in llm.stream_complete(message):
+    #     print(chunk.delta, end="", flush=True)
+    #     response += str(chunk.delta)
+    #     yield response
+    outputs = []
     for chunk in llm.stream_complete(message):
-        print(chunk.delta, end="", flush=True)
-        response += str(chunk.delta)
-        yield response
+        outputs.append(chunk.delta)
+        if chunk.delta in stop_tokens:
+            break
+        yield "".join(outputs)
 
 demo = gr.ChatInterface(
     respond,