streaming per token

app.py

@@ -29,10 +29,10 @@ pipe = Llama(
 def predict(message: str, history: List[List[str]], max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS):
     if not message:
         return "", history
-    
+
     prompt = message
     history.append([message, ""])
-    
+
     # Initialize reply for this round
     reply = ""
     
@@ -44,12 +44,13 @@ def predict(message: str, history: List[List[str]], max_new_tokens: int = DEFAUL
         stream=True
     )
     
+    # Send each token stream output to the user
     for output in stream:
         new_text = output['choices'][0]['text']
-        reply += new_text 
-        history[-1][1] = reply  # Update the current reply in history
-        yield "", history 
-    return "", history  # Always return at the end to terminate the generator
+        reply += new_text
+        history[-1][1] = reply # Update the current reply in history
+        yield "", history
+
 
 with gr.Blocks() as demo:
     gr.Markdown(DESCRIPTION)
@@ -64,4 +65,4 @@ with gr.Blocks() as demo:
     )
     textbox.submit(predict, [textbox, chatbot, max_new_tokens_slider], [textbox, chatbot])
 
-demo.queue().launch()
+demo.queue().launch()