Update app.py

app.py

@@ -1,4 +1,5 @@
 import math
+import time
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -258,7 +259,7 @@ model, tokenizer, load_status = load_model()
 
 @torch.no_grad()
 def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):
-    """Autoregressive generation with ROSA updates, yielding generated token lists step-by-step."""
+    """Autoregressive generation with ROSA updates, yielding generated token lists and speed (tok/s) step-by-step."""
     device = next(model.parameters()).device
     eos_id = tokenizer.eos_token_id
 
@@ -270,6 +271,8 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
     logits = out.logits[0, -1]
     generated = list(ids)
     reply_tokens = []
+    
+    start_time = time.perf_counter()
 
     for _ in range(max_new_tokens):
         scaled = logits / max(temperature, 1e-5)
@@ -283,7 +286,9 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
             break
 
         reply_tokens.append(next_token)
-        yield reply_tokens
+        elapsed = time.perf_counter() - start_time
+        tps = len(reply_tokens) / elapsed if elapsed > 0 else 0.0
+        yield reply_tokens, tps
 
         # ROSA for the extended sequence, use the last prediction
         next_rosa = rosa(generated)[-1]
@@ -296,7 +301,7 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
 
 
 def chat_function(message, history):
-    """Gradio ChatInterface streaming callback."""
+    """Gradio ChatInterface streaming callback with tok/s throughput display."""
     messages = []
     for turn in history:
         if isinstance(turn, (list, tuple)):
@@ -330,10 +335,10 @@ def chat_function(message, history):
     # Cue the assistant to start replying
     ids.append(asst_id)
 
-    # Stream generated tokens in real time
-    for reply_tokens in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50):
+    # Stream generated tokens and tok/s throughput in real time
+    for reply_tokens, tps in generate_reply_stream(ids, max_new_tokens=150, temperature=0.8, top_k=50):
         reply = tokenizer.decode(reply_tokens, skip_special_tokens=True).strip()
-        yield reply
+        yield f"{reply}\n\n⚡ *{tps:.1f} tok/s*"
 
 
 # -------------------------------------------------------------------------