Update app.py

app.py

@@ -12,11 +12,8 @@ from transformers import (
 )
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
-# -------------------------------------------------------------------------
-# Model components (matches saved checkpoint architecture)
-# -------------------------------------------------------------------------
-
 class FWKVConfig(PretrainedConfig):
+    """Configuration class for FWKV-ROSA model architecture."""
     model_type = "fwkv"
 
     def __init__(
@@ -40,7 +37,6 @@ class FWKVConfig(PretrainedConfig):
         self.seq_len = seq_len
         self.wkv_floor = wkv_floor
 
-
 def rosa(x: list[int]) -> list[int]:
     """Causal copy‑signal predictor; returns y[i] = token after longest
     repeating suffix ending at i, or -1 if none."""
@@ -99,7 +95,6 @@ def rosa(x: list[int]) -> list[int]:
             v = link[v]
     return y
 
-
 def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
     """Hillis–Steele inclusive scan with constant per‑channel decay."""
     W = W.to(dtype=a.dtype)          # keep precision
@@ -112,8 +107,8 @@ def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
         d *= 2
     return val
 
-
 class FactorizedTiedHead(nn.Module):
+    """Factorized embedding projection and tied output head."""
     def __init__(self, vocab_size: int, d_model: int, d_emb: int):
         super().__init__()
         self.d_model = d_model
@@ -130,8 +125,8 @@ class FactorizedTiedHead(nn.Module):
     def logits(self, x_emb):
         return F.linear(x_emb, self.weight)
 
-
 class FWKVBlock(nn.Module):
+    """FWKV layer with linear time attention-style recurrent mechanism."""
     def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):
         super().__init__()
         self.floor = floor
@@ -171,8 +166,8 @@ class FWKVBlock(nn.Module):
         x = self.norm_ffn(x + self.ffn(x))
         return x, new_state
 
-
 class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
+    """Full causal language model utilizing FWKV recurrent layers and ROSA embeddings."""
     config_class = FWKVConfig
 
     def __init__(self, config):
@@ -231,11 +226,6 @@ class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
         return {"input_ids": input_ids, "rosa_ids": rosa_ids,
                 "past_key_values": past_key_values, "use_cache": True}
 
-
-# -------------------------------------------------------------------------
-# Loading & Streaming Helpers
-# -------------------------------------------------------------------------
-
 USER_TOKEN = "<|user|>"
 ASSISTANT_TOKEN = "<|assistant|>"
 
@@ -259,11 +249,11 @@ 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 and speed (tok/s) step-by-step."""
+    """Autoregressive generation with ROSA updates, yielding token lists and comprehensive throughput metrics."""
     device = next(model.parameters()).device
     eos_id = tokenizer.eos_token_id
 
-    # initial forward over the whole prompt
+    # Initial forward pass over the prompt
     inp = torch.tensor([ids], device=device)
     rosa_ids = torch.tensor([rosa(ids)], device=device)
     out = model(input_ids=inp, rosa_ids=rosa_ids, use_cache=True)
@@ -273,6 +263,8 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
     reply_tokens = []
     
     start_time = time.perf_counter()
+    prev_step_time = start_time
+    instant_tps_list = []
 
     for _ in range(max_new_tokens):
         scaled = logits / max(temperature, 1e-5)
@@ -286,11 +278,33 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
             break
 
         reply_tokens.append(next_token)
-        elapsed = time.perf_counter() - start_time
-        tps = len(reply_tokens) / elapsed if elapsed > 0 else 0.0
-        yield reply_tokens, tps
+        now = time.perf_counter()
+
+        # Calculate per-step instant duration and speed
+        step_duration = now - prev_step_time
+        prev_step_time = now
 
-        # ROSA for the extended sequence, use the last prediction
+        if step_duration > 0:
+            instant_tps = 1.0 / step_duration
+            instant_tps_list.append(instant_tps)
+
+        # Compute aggregate throughput metrics
+        total_elapsed = now - start_time
+        avg_tps = len(reply_tokens) / total_elapsed if total_elapsed > 0 else 0.0
+        current_tps = instant_tps_list[-1] if instant_tps_list else avg_tps
+        min_tps = min(instant_tps_list) if instant_tps_list else avg_tps
+        max_tps = max(instant_tps_list) if instant_tps_list else avg_tps
+
+        stats = {
+            "current": current_tps,
+            "avg": avg_tps,
+            "min": min_tps,
+            "max": max_tps,
+        }
+
+        yield reply_tokens, stats
+
+        # ROSA prediction for the next step
         next_rosa = rosa(generated)[-1]
         step_inp = torch.tensor([[next_token]], device=device)
         step_rosa = torch.tensor([[next_rosa]], device=device)
@@ -299,9 +313,8 @@ def generate_reply_stream(ids: list[int], max_new_tokens=150, temperature=0.8, t
         states = out.past_key_values
         logits = out.logits[0, -1]
 
-
 def chat_function(message, history):
-    """Gradio ChatInterface streaming callback with tok/s throughput display."""
+    """Gradio ChatInterface streaming handler formatted with speed stats (Live, Avg, Min, Max)."""
     messages = []
     for turn in history:
         if isinstance(turn, (list, tuple)):
@@ -313,7 +326,7 @@ def chat_function(message, history):
             messages.append(turn)
     messages.append({"role": "user", "content": message})
 
-    # Convert to token IDs
+    # Encode token sequence according to model chat template
     user_id = tokenizer.convert_tokens_to_ids(USER_TOKEN)
     asst_id = tokenizer.convert_tokens_to_ids(ASSISTANT_TOKEN)
     eos_id = tokenizer.eos_token_id
@@ -327,23 +340,22 @@ def chat_function(message, history):
         elif role == "assistant":
             ids += [asst_id] + content_ids + [eos_id]
 
-    # Truncate from the left if needed
+    # Truncate left if context exceeds model max sequence length
     max_len = model.config.seq_len if model else 1024
     if len(ids) > max_len:
         ids = ids[-max_len:]
 
-    # Cue the assistant to start replying
+    # Prompt assistant response
     ids.append(asst_id)
 
-    # 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):
+    # Stream generated output with full throughput statistics
+    for reply_tokens, stats 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 f"{reply}\n\n⚡ *{tps:.1f} tok/s*"
-
-
-# -------------------------------------------------------------------------
-# Gradio UI
-# -------------------------------------------------------------------------
+        metrics_bar = (
+            f"⚡ **{stats['current']:.1f} tok/s** "
+            f"*(Avg: **{stats['avg']:.1f}** | Min: **{stats['min']:.1f}** | Max: **{stats['max']:.1f}** tok/s)*"
+        )
+        yield f"{reply}\n\n{metrics_bar}"
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"""