tps field

app.py

@@ -40,6 +40,9 @@ def predict(message: str, history: List[List[str]], max_new_tokens: int = DEFAUL
     # Initialize token count and start time
     token_count = 0
     start_time = time.time()
+    
+    last_token_count = 0
+    last_time = start_time
 
     # This will produce a generator of output chunks
     stream = pipe(
@@ -55,16 +58,22 @@ def predict(message: str, history: List[List[str]], max_new_tokens: int = DEFAUL
         reply += new_text
         token_count += len(new_text.split()) # Estimate tokens by counting spaces
         history[-1][1] = reply # Update the current reply in history
-    
-        # Calculate elapsed time and TPS
-        elapsed_time = time.time() - start_time
+        
+        # Calculate elapsed time since last update
+        elapsed_time = time.time() - last_time
+        
         if elapsed_time > 0:
-            tps = token_count / elapsed_time
+            # Calculate tokens per second since last update
+            tokens_per_second = (token_count - last_token_count) / elapsed_time
         else:
-            tps = 0
-        
+            tokens_per_second = 0
+
         # Update the status using gradio's progress
-        progress(message=f"Tokens per second: {tps:.2f}")
+        progress(message=f"Tokens per second: {tokens_per_second:.2f}")
+
+        # Update for next iteration
+        last_token_count = token_count
+        last_time = time.time()
 
         yield "", history
 
@@ -80,7 +89,7 @@ with gr.Blocks() as demo:
         value=DEFAULT_MAX_NEW_TOKENS,
         label="Max New Tokens",
     )
-    status_field = gr.Text(label="Status", interactive=False, visible=True) # Add Status field
+    
     textbox.submit(predict, [textbox, chatbot, max_new_tokens_slider], [textbox, chatbot], )
 
 demo.queue().launch()