updated and cleaned up the streaming code (#15)

- updated and cleaned up the streaming code (63c0b01fef5a1fddd405ffccfe51a81e493ca700)


Co-authored-by: yuvraj sharma <[email protected]>

app.py

@@ -29,18 +29,18 @@ class StopOnTokens(StoppingCriteria):
                 return True
         return False
 
-def user(user_message, history):
+def chat(curr_system_message, user_message, history):
+    # Append the user's message to the conversation history
     history = history + [[user_message, ""]]
-    return "", history, history
-
-
-def bot(history, curr_system_message):
+    # Initialize a StopOnTokens object 
     stop = StopOnTokens()
+
+    # Construct the input message string for the model by concatenating the current system message and conversation history
     messages = curr_system_message + \
         "".join(["".join(["<|USER|>"+item[0], "<|ASSISTANT|>"+item[1]])
                 for item in history])
     
-    #model_inputs = tok([messages], return_tensors="pt")['input_ids'].cuda()[:, :4096-1024]
+    # Tokenize the messages string
     model_inputs = tok([messages], return_tensors="pt").to("cuda")
     streamer = TextIteratorStreamer(tok, timeout=10., skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
@@ -58,16 +58,18 @@ def bot(history, curr_system_message):
     t.start()
 
     print(history)
+    # Initialize an empty string to store the generated text
+    partial_text = ""
     for new_text in streamer:
         print(new_text)
-        history[-1][1] += new_text
-        yield history, history
-        
-    return history, history
+        partial_text += new_text
+        history[-1][1] = partial_text
+        # Yield an empty string to cleanup the message textbox and the updated conversation history
+        yield "", history
 
 
 with gr.Blocks() as demo:
-    history = gr.State([])
+    #history = gr.State([])
     gr.Markdown("## StableLM-Tuned-Alpha-7b Chat")
     gr.HTML('''<center><a href="https://huggingface.co/spaces/stabilityai/stablelm-tuned-alpha-chat?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to skip the queue and run in a private space</center>''')
     chatbot = gr.Chatbot().style(height=500)
@@ -81,10 +83,9 @@ with gr.Blocks() as demo:
     system_msg = gr.Textbox(
         start_message, label="System Message", interactive=False, visible=False)
 
-    msg.submit(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then(
-        fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True)
-    submit.click(fn=user, inputs=[msg, history], outputs=[msg, chatbot, history], queue=False).then(
-        fn=bot, inputs=[chatbot, system_msg], outputs=[chatbot, history], queue=True)
-    clear.click(lambda: [None, []], None, [chatbot, history], queue=False)
+    msg.submit(fn=chat, inputs=[system_msg, msg, chatbot], outputs=[msg, chatbot], queue=True)
+    submit.click(fn=chat, inputs=[system_msg, msg, chatbot], outputs=[msg, chatbot], queue=True)
+    clear.click(lambda: [None, []], None, [chatbot], queue=False)
+
 demo.queue(concurrency_count=2)
 demo.launch()