Update app.py

app.py

@@ -12,7 +12,7 @@ from transformers import (
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 # -------------------------------------------------------------------------
-# Model components (must match the saved checkpoint exactly)
+# Model components (matches saved checkpoint architecture)
 # -------------------------------------------------------------------------
 
 class FWKVConfig(PretrainedConfig):
@@ -40,7 +40,6 @@ class FWKVConfig(PretrainedConfig):
         self.wkv_floor = wkv_floor
 
 
-# ── ROSA (exact copy from training) ─────────────────────────────────────
 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."""
@@ -100,7 +99,6 @@ def rosa(x: list[int]) -> list[int]:
     return y
 
 
-# ── Vectorised parallel scan ─────────────────────────────────────────────
 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
@@ -114,7 +112,6 @@ def parallel_scan_decay(a: torch.Tensor, W: torch.Tensor) -> torch.Tensor:
     return val
 
 
-# ── Factorised tied head ─────────────────────────────────────────────────
 class FactorizedTiedHead(nn.Module):
     def __init__(self, vocab_size: int, d_model: int, d_emb: int):
         super().__init__()
@@ -133,7 +130,6 @@ class FactorizedTiedHead(nn.Module):
         return F.linear(x_emb, self.weight)
 
 
-# ── FWKV Block (with parallel scan) ──────────────────────────────────────
 class FWKVBlock(nn.Module):
     def __init__(self, d: int, ffn_mult: int = 4, floor: float = 0.1):
         super().__init__()
@@ -175,7 +171,6 @@ class FWKVBlock(nn.Module):
         return x, new_state
 
 
-# ── Full Language Model ──────────────────────────────────────────────────
 class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
     config_class = FWKVConfig
 
@@ -203,7 +198,6 @@ class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
         **kwargs,
     ):
         if rosa_ids is None:
-            # fallback: compute on the fly (expensive, but safe)
             rows = [rosa(row.tolist()) for row in input_ids.detach().cpu()]
             rosa_ids = torch.tensor(rows, device=input_ids.device, dtype=torch.long)
 
@@ -238,7 +232,7 @@ class FWKVLanguageModel(PreTrainedModel, GenerationMixin):
 
 
 # -------------------------------------------------------------------------
-# Loading & Chat helpers
+# Loading & Streaming Helpers
 # -------------------------------------------------------------------------
 
 USER_TOKEN = "<|user|>"
@@ -263,8 +257,8 @@ def load_model():
 model, tokenizer, load_status = load_model()
 
 @torch.no_grad()
-def generate_reply(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50):
-    """Autoregressive generation with ROSA updates, returning full sequence."""
+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."""
     device = next(model.parameters()).device
     eos_id = tokenizer.eos_token_id
 
@@ -275,6 +269,7 @@ def generate_reply(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50
     states = out.past_key_values
     logits = out.logits[0, -1]
     generated = list(ids)
+    reply_tokens = []
 
     for _ in range(max_new_tokens):
         scaled = logits / max(temperature, 1e-5)
@@ -287,6 +282,9 @@ def generate_reply(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50
         if next_token == eos_id:
             break
 
+        reply_tokens.append(next_token)
+        yield reply_tokens
+
         # ROSA for the extended sequence, use the last prediction
         next_rosa = rosa(generated)[-1]
         step_inp = torch.tensor([[next_token]], device=device)
@@ -296,16 +294,18 @@ def generate_reply(ids: list[int], max_new_tokens=150, temperature=0.8, top_k=50
         states = out.past_key_values
         logits = out.logits[0, -1]
 
-    return generated
-
 
 def chat_function(message, history):
-    """Gradio ChatInterface callback. history is a list of (user, assistant) pairs."""
-    # Build the full conversation in the chat template
+    """Gradio ChatInterface streaming callback."""
     messages = []
-    for user_msg, asst_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": asst_msg})
+    for turn in history:
+        if isinstance(turn, (list, tuple)):
+            user_msg, asst_msg = turn
+            messages.append({"role": "user", "content": user_msg})
+            if asst_msg:
+                messages.append({"role": "assistant", "content": asst_msg})
+        elif isinstance(turn, dict):
+            messages.append(turn)
     messages.append({"role": "user", "content": message})
 
     # Convert to token IDs
@@ -323,25 +323,23 @@ def chat_function(message, history):
             ids += [asst_id] + content_ids + [eos_id]
 
     # Truncate from the left if needed
-    max_len = model.config.seq_len
+    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
     ids.append(asst_id)
 
-    # Generate the assistant's reply
-    full_gen = generate_reply(ids, max_new_tokens=150, temperature=0.8, top_k=50)
+    # Stream generated tokens in real time
+    for reply_tokens 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
 
-    # Extract only the new assistant tokens (after the last asst_id)
-    # Find the position of the last asst_id and take everything after it
-    assistant_start = len(ids) - 1  # index of the asst_id we just appended
-    reply_ids = full_gen[assistant_start + 1:]  # skip the asst_id itself
-    reply = tokenizer.decode(reply_ids, skip_special_tokens=True).strip()
-    return reply
 
+# -------------------------------------------------------------------------
+# Gradio UI
+# -------------------------------------------------------------------------
 
-# ── Gradio UI ────────────────────────────────────────────────────────────
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f"""
     # ⚡ FWKV-ROSA Chat