Update app.py

app.py

@@ -53,6 +53,7 @@ def chat(message, history, temperature, max_new_tokens):
     
     # Stream the response
     start_time = time.time()
+    token_count = 0
     
     # Create a TextStreamer for token streaming
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
@@ -73,9 +74,18 @@ def chat(message, history, temperature, max_new_tokens):
     outputs = []
     for new_token in streamer:
         outputs.append(new_token)
+        token_count += 1
+        
+        # Calculate tokens per second
+        elapsed_time = time.time() - start_time
+        tokens_per_second = token_count / elapsed_time if elapsed_time > 0 else 0
+        
+        # Update the token status
+        token_status_value = f"Tokens Generated: {token_count}, Tokens/Second: {tokens_per_second:.2f}"
+        yield "".join(outputs), token_status_value
+        
         if any(stop_token in new_token for stop_token in stop_tokens):
             break
-        yield "".join(outputs)
 
 def reload_model_button():
     """Reload the model manually via a button."""
@@ -114,8 +124,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     def respond(message, chat_history, temperature, max_new_tokens):
         bot_message = ""
-        for partial_response in chat(message, chat_history, temperature, max_new_tokens):
+        for partial_response, token_status_value in chat(message, chat_history, temperature, max_new_tokens):
             bot_message = partial_response
+            token_status.update(value=token_status_value)  # Update token generation status
             yield "", chat_history + [(message, bot_message)]
     
     send_button.click(respond, inputs=[textbox, chatbot, temperature_slider, max_tokens_slider], outputs=[textbox, chatbot])