Update app.py

app.py

@@ -2,12 +2,11 @@ import os
 import re
 import json
 import time
-import shutil
 import torch
 import torch.nn as nn
 import gradio as gr
 from pathlib import Path
-from collections import deque  # ADDED for penalty window
+from collections import deque
 
 # ─────────────────────────────────────────
 # 🔐 Admin password
@@ -23,7 +22,6 @@ CONFIG_PATH = SPACE_ROOT / "config.json"
 _pt_files  = sorted(SPACE_ROOT.glob("*.pt"))
 MODEL_PATH = _pt_files[0] if _pt_files else SPACE_ROOT / "default_model.pt"
 
-# Find tokenizer — any .json that isn't config.json
 _tok_files    = [f for f in SPACE_ROOT.glob("*.json") if f.name != "config.json"]
 TOKENIZER_PATH = _tok_files[0] if _tok_files else None
 
@@ -31,7 +29,7 @@ TOKENIZER_PATH = _tok_files[0] if _tok_files else None
 # 🗃️ Config
 # ─────────────────────────────────────────
 DEFAULT_CONFIG = {
-    "model_type":      "char",   # "char" or "token"
+    "model_type":      "char",
     "hidden_layers":   5,
     "neurons":         768,
     "embed_size":      384,
@@ -42,17 +40,16 @@ DEFAULT_CONFIG = {
     "bot_tag":         "### Response:",
     "eos_token":       "<|end|>",
     "system_prompt":   "You are a helpful and intelligent AI assistant named Linny.",
-    "default_temp":    0.7,
-    "default_penalty": 1.10,  # CHANGED: from 1.2 to 1.10 (better default)
-    "default_penalty_window": 110,  # ADDED
-    "default_top_p":   0.4,    # CHANGED: from 0.9 to 0.4 (matches local)
-    "default_top_k":   65,     # CHANGED: from 50 to 65 (matches local)
-    "default_max_len": 1575,
-    "min_response_tokens": 3,   # ADDED
-    "max_reasoning_tokens": 2500,  # ADDED
-    # reasoning_mode: "prompt_suffix" appends /think | "response_prefix" prefills <think>
-    "reasoning_mode":  "prompt_suffix",
-    "reasoning_start": True,
+    "default_temp":    0.8,
+    "default_penalty": 1.05,
+    "default_penalty_window": 110,
+    "default_top_p":   0.4,
+    "default_top_k":   65,
+    "default_max_len": 2600,
+    "reasoning_mode":  "response_prefix",
+    "reasoning_start": False,
+    "min_response_tokens": 3,
+    "max_reasoning_tokens": 2500,
 }
 
 def load_config() -> dict:
@@ -68,7 +65,6 @@ def save_config(cfg: dict):
     with open(CONFIG_PATH, "w") as f:
         json.dump(cfg, f, indent=2)
 
-
 # ─────────────────────────────────────────
 # 🧠 Model Architectures (unchanged)
 # ─────────────────────────────────────────
@@ -77,32 +73,25 @@ class LSTMCharLM(nn.Module):
         super().__init__()
         self.embed = nn.Embedding(vocab_size, embed_size)
         self.lstm  = nn.LSTM(embed_size, hidden_size, num_layers=num_layers,
-                             batch_first=True,
-                             dropout=dropout if num_layers > 1 else 0)
-        self.fc = nn.Linear(hidden_size, vocab_size)
-
+                             batch_first=True, dropout=dropout if num_layers > 1 else 0)
+        self.fc    = nn.Linear(hidden_size, vocab_size)
     def forward(self, x, hidden=None):
-        e = self.embed(x)
-        out, hidden = self.lstm(e, hidden)
+        out, hidden = self.lstm(self.embed(x), hidden)
         return self.fc(out), hidden
 
-
 class LSTMTokenLM(nn.Module):
     def __init__(self, vocab_size, embed_size, hidden_size, num_layers, dropout=0.2):
         super().__init__()
         self.embed = nn.Embedding(vocab_size, embed_size)
         self.lstm  = nn.LSTM(embed_size, hidden_size, num_layers=num_layers,
-                             batch_first=True,
-                             dropout=dropout if num_layers > 1 else 0)
-        self.fc = nn.Linear(hidden_size, vocab_size)
-
+                             batch_first=True, dropout=dropout if num_layers > 1 else 0)
+        self.fc    = nn.Linear(hidden_size, vocab_size)
     def forward(self, x, hidden=None):
         out, hidden = self.lstm(self.embed(x), hidden)
         return self.fc(out), hidden
 
-
 # ─────────────────────────────────────────
-# 🔤 GPT-2 byte decoder (token mode only)
+# 🔤 GPT-2 byte decoder
 # ─────────────────────────────────────────
 def _build_byte_decoder():
     bs = (list(range(ord('!'), ord('~')+1)) +
@@ -113,7 +102,7 @@ def _build_byte_decoder():
     for b in range(256):
         if b not in bs:
             bs.append(b)
-            cs.append(256 + n)
+            cs.append(256+n)
             n += 1
     return {chr(c): b for b, c in zip(bs, cs)}
 
@@ -125,122 +114,60 @@ def _tok_to_bytes(tok_str):
     except KeyError:
         return tok_str.encode('utf-8', errors='replace')
 
-
 # ─────────────────────────────────────────
-# ⚙️ Model Loader (UPDATED with caching and optimizations)
+# ⚙️ Model Loader (identical generation to local)
 # ─────────────────────────────────────────
 class LinnyModel:
     def __init__(self, pt_path, config: dict, tokenizer_path=None):
         self.config     = config
         self.model_type = config.get("model_type", "char")
         self.device     = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # ADDED: Cache for hidden states to speed up repeated prefixes
-        self.cached_hidden = None
-        self.cached_prefix = None
-        self.cached_prefix_ids = None
 
         ckpt = torch.load(pt_path, map_location=self.device, weights_only=False)
 
         if self.model_type == "token":
-            try:
-                from tokenizers import Tokenizer as HFTokenizer
-            except ImportError:
-                raise ImportError("Install tokenizers: pip install tokenizers")
-
+            from tokenizers import Tokenizer as HFTokenizer
             tok_path = tokenizer_path or config.get("tokenizer_path")
             if not tok_path or not Path(str(tok_path)).exists():
                 raise FileNotFoundError(f"Tokenizer not found: {tok_path}")
-
             self.tokenizer = HFTokenizer.from_file(str(tok_path))
-            vocab_size     = self.tokenizer.get_vocab_size()
-            self.chars     = None
-            self.stoi      = None
-            self.itos      = None
-
-            arch    = ckpt.get('config', {})
+            vocab_size = self.tokenizer.get_vocab_size()
+            self.chars = None; self.stoi = None; self.itos = None
+            arch = ckpt.get('config', {})
             layers  = arch.get('hidden_layers', config['hidden_layers'])
             neurons = arch.get('neurons',       config['neurons'])
             embed   = arch.get('embed_size',    config['embed_size'])
             dropout = arch.get('dropout',       config.get('dropout', 0.2))
-
             self.model = LSTMTokenLM(vocab_size, embed, neurons, layers, dropout).to(self.device)
             self.model.load_state_dict(ckpt['model_state'])
-
         else:
             self.tokenizer = None
-            self.chars     = ckpt["chars"]
-            self.stoi      = {ch: i for i, ch in enumerate(self.chars)}
-            self.itos      = {i: ch for i, ch in enumerate(self.chars)}
-
+            self.chars = ckpt["chars"]
+            self.stoi = {ch: i for i, ch in enumerate(self.chars)}
+            self.itos = {i: ch for i, ch in enumerate(self.chars)}
             dropout = ckpt.get("config", {}).get("dropout", config.get("dropout", 0.2))
-            self.model = LSTMCharLM(
-                len(self.chars),
-                config["embed_size"],
-                config["neurons"],
-                config["hidden_layers"],
-                dropout,
-            ).to(self.device)
+            self.model = LSTMCharLM(len(self.chars), config["embed_size"],
+                                    config["neurons"], config["hidden_layers"], dropout).to(self.device)
             self.model.load_state_dict(ckpt["model_state"])
 
-        # ADDED: Try to use torch.compile for speed (PyTorch 2.0+)
-        if hasattr(torch, 'compile') and self.device == torch.device("cpu"):
-            try:
-                self.model = torch.compile(self.model, mode="reduce-overhead")
-                print("✅ TorchScript compilation enabled for CPU speedup")
-            except:
-                print("⚠️ Torch.compile failed, using regular model")
-        
         self.model.eval()
         self.epoch = ckpt.get('epoch', '?')
 
-    def _prime_hidden(self, text, hidden):
-        if self.model_type == "token":
-            ids = self.tokenizer.encode(text).ids
-        else:
-            ids = [self.stoi.get(c, 0) for c in text]
-        t = torch.tensor([ids], dtype=torch.long).to(self.device)
-        with torch.no_grad():
-            _, hidden = self.model(t, hidden)
-        return hidden, torch.tensor([[ids[-1]]], dtype=torch.long).to(self.device)
-    
-    # ADDED: Method to get cached hidden state for repeated prompts
-    def get_cached_prefix(self, prefix_text):
-        """Cache hidden state for a prefix to avoid recomputation"""
-        if self.cached_prefix == prefix_text:
-            return self.cached_hidden, self.cached_prefix_ids
-        
-        if self.model_type == "token":
-            ids = self.tokenizer.encode(prefix_text).ids
-        else:
-            ids = [self.stoi.get(c, 0) for c in prefix_text]
-        
-        if not ids:
-            return None, None
-            
-        t = torch.tensor([ids], dtype=torch.long).to(self.device)
-        with torch.no_grad():
-            _, hidden = self.model(t, None)
-        
-        self.cached_prefix = prefix_text
-        self.cached_hidden = hidden
-        self.cached_prefix_ids = torch.tensor([[ids[-1]]], dtype=torch.long).to(self.device)
-        
-        return hidden, self.cached_prefix_ids
-
-    def stream_generate(self, prompt, temperature=0.7, max_len=1575,
-                        penalty=1.2, penalty_window=110,  # ADDED penalty_window
-                        top_p=0.9, top_k=50, force_thinking=False,
-                        prefix_text="", penalize_prefix=False,  # ADDED for continue feature
-                        min_response_tokens=3,  # ADDED
-                        max_reasoning_tokens=2500):  # ADDED
+    # ------------------------------------------------------------------
+    # Exact copy of local server's generate_stream logic
+    # ------------------------------------------------------------------
+    def stream_generate(self, prompt, temperature=0.8, max_len=2600,
+                        penalty=1.05, penalty_window=110,
+                        top_p=0.4, top_k=65, force_thinking=False,
+                        prefix_text="", penalize_prefix=False,
+                        min_response_tokens=3, max_reasoning_tokens=2500):
         """
-        Generate with proper repetition penalty, reasoning tracking, and continue support.
+        Matches local linny_server.py generate_stream exactly.
         """
-        cfg      = self.config
+        cfg = self.config
         user_tag = cfg.get("user_tag", "### Instruction:")
-        bot_tag  = cfg.get("bot_tag",  "### Response:")
-        r_mode   = cfg.get("reasoning_mode", "prompt_suffix")
+        bot_tag  = cfg.get("bot_tag", "### Response:")
+        r_mode   = cfg.get("reasoning_mode", "response_prefix")
         eos_token_str = cfg.get("eos_token", "<|end|>")
 
         # Apply prompt_suffix mode
@@ -250,32 +177,22 @@ class LinnyModel:
                 actual_prompt = prompt.strip() + " /think"
 
         formatted = f"{user_tag}\n{actual_prompt}\n\n{bot_tag}\n"
-        
+
+        # Working memory cap
         working_memory = cfg.get("working_memory", 0)
         if working_memory > 0:
             max_len = max(50, min(max_len, working_memory - len(formatted)))
 
         hidden = None
         generated = ""
-        
-        # ADDED: Track recent tokens for penalty
         recent_tokens = deque(maxlen=penalty_window)
-        
-        # ADDED: Track reasoning state
+
+        # State tracking (exactly as local)
         in_reasoning = False
         think_closed = False
         awaiting_response = False
         response_token_count = 0
         reasoning_toks = 0
-        
-        # Get token IDs for special tokens (token mode only)
-        eos_id = None
-        think_open_id = None
-        think_close_id = None
-        if self.model_type == "token":
-            eos_id = self.tokenizer.token_to_id(eos_token_str)
-            think_open_id = self.tokenizer.token_to_id("<think>")
-            think_close_id = self.tokenizer.token_to_id("</think>")
 
         with torch.no_grad():
             # Encode the conversation prefix
@@ -284,44 +201,47 @@ class LinnyModel:
             else:
                 ids = [self.stoi.get(c, 0) for c in formatted]
 
-            t = torch.tensor([ids], dtype=torch.long).to(self.device)
+            t = torch.tensor([ids], dtype=torch.long, device=self.device)
             _, hidden = self.model(t, hidden)
-            input_token = torch.tensor([[ids[-1]]], dtype=torch.long).to(self.device)
-            
-            # If we have existing assistant response (prefix_text for continue mode)
+            input_token = torch.tensor([[ids[-1]]], dtype=torch.long, device=self.device)
+
+            # If we have existing assistant response (continue mode)
             if prefix_text:
                 if self.model_type == "token":
                     prefix_ids = self.tokenizer.encode(prefix_text).ids
                 else:
                     prefix_ids = [self.stoi.get(c, 0) for c in prefix_text]
-                
                 if prefix_ids:
-                    pt = torch.tensor([prefix_ids], dtype=torch.long).to(self.device)
+                    pt = torch.tensor([prefix_ids], dtype=torch.long, device=self.device)
                     _, hidden = self.model(pt, hidden)
                     generated = prefix_text
-                    input_token = torch.tensor([[prefix_ids[-1]]], dtype=torch.long).to(self.device)
+                    input_token = torch.tensor([[prefix_ids[-1]]], dtype=torch.long, device=self.device)
                     if penalize_prefix:
                         recent_tokens.extend(prefix_ids)
-                    
                     # Update state based on prefix
                     if "<think>" in prefix_text and "</think>" not in prefix_text:
                         in_reasoning = True
-                        think_closed = False
                     elif "</think>" in prefix_text:
                         think_closed = True
                         awaiting_response = True
 
-            # Prefill <think> if in response_prefix mode (only if no prefix)
+            # Prefill <think> if in response_prefix mode and no prefix
             if not prefix_text and force_thinking and r_mode == "response_prefix":
-                if self.model_type == "token" and think_open_id is not None:
-                    tt = torch.tensor([[think_open_id]], dtype=torch.long).to(self.device)
+                if self.model_type == "token":
+                    think_id = self.tokenizer.token_to_id("<think>")
+                    if think_id is not None:
+                        tt = torch.tensor([[think_id]], dtype=torch.long, device=self.device)
+                        _, hidden = self.model(tt, hidden)
+                        input_token = tt
+                        generated = "<think>"
+                        yield "<think>"
+                        in_reasoning = True
+                else:
+                    # Char mode: prime hidden with "<think>"
+                    think_ids = [self.stoi.get(ch, 0) for ch in "<think>"]
+                    tt = torch.tensor([think_ids], dtype=torch.long, device=self.device)
                     _, hidden = self.model(tt, hidden)
-                    input_token = tt
-                    generated = "<think>"
-                    yield "<think>"
-                    in_reasoning = True
-                elif self.model_type == "char":
-                    hidden, input_token = self._prime_hidden("<think>", hidden)
+                    input_token = torch.tensor([[think_ids[-1]]], dtype=torch.long, device=self.device)
                     generated = "<think>"
                     for ch in "<think>":
                         yield ch
@@ -330,25 +250,40 @@ class LinnyModel:
                 if cfg.get("reasoning_start", False) and not prefix_text:
                     prefix = f"I need to think about this. The user said '{prompt}'"
                     if self.model_type == "token":
-                        hidden, input_token = self._prime_hidden(prefix, hidden)
+                        pids = self.tokenizer.encode(prefix).ids
+                        pt = torch.tensor([pids], dtype=torch.long, device=self.device)
+                        _, hidden = self.model(pt, hidden)
+                        input_token = torch.tensor([[pids[-1]]], dtype=torch.long, device=self.device)
                         generated += prefix
                         yield prefix
                     else:
-                        hidden, input_token = self._prime_hidden(prefix, hidden)
-                        generated += prefix
                         for ch in prefix:
+                            idx = self.stoi.get(ch, 0)
+                            it = torch.tensor([[idx]], dtype=torch.long, device=self.device)
+                            _, hidden = self.model(it, hidden)
+                            input_token = it
+                            generated += ch
                             yield ch
 
-            # ── Token generation ──
+            # Get special token IDs
+            eos_id = None
+            think_open_id = None
+            think_close_id = None
+            if self.model_type == "token":
+                eos_id = self.tokenizer.token_to_id(eos_token_str)
+                think_open_id = self.tokenizer.token_to_id("<think>")
+                think_close_id = self.tokenizer.token_to_id("</think>")
+
+            # ------------------------------------------------------------------
+            # Token generation (exactly as local)
+            # ------------------------------------------------------------------
             if self.model_type == "token":
                 byte_buf = b""
-                token_count = 0
-                
                 for step in range(max_len):
                     logits, hidden = self.model(input_token, hidden)
                     lf = logits[0, -1].float() / max(temperature, 1e-8)
 
-                    # FIXED: Proper repetition penalty
+                    # Repetition penalty
                     if penalty != 1.0 and len(recent_tokens) > 0:
                         penalized_ids = set(recent_tokens)
                         for token_id in penalized_ids:
@@ -369,16 +304,14 @@ class LinnyModel:
                         lf[si[rm]] = float("-inf")
 
                     nxt = torch.multinomial(torch.softmax(lf, dim=-1), 1).item()
-                    token_count += 1
 
                     # EOS handling with forced minimum response
-                    if eos_id is not None and nxt == eos_id:
+                    if nxt == eos_id:
                         if awaiting_response and response_token_count < min_response_tokens:
-                            continue  # Not enough response yet, skip EOS
+                            continue
                         else:
                             break
 
-                    # Update recent tokens for penalty
                     recent_tokens.append(nxt)
 
                     # Update reasoning state
@@ -390,24 +323,22 @@ class LinnyModel:
                         think_closed = True
                         awaiting_response = True
                         response_token_count = 0
-                    
+
                     if in_reasoning and not think_closed:
                         reasoning_toks += 1
-                        if max_reasoning_tokens and max_reasoning_tokens > 0 and reasoning_toks >= max_reasoning_tokens:
+                        if max_reasoning_tokens and reasoning_toks >= max_reasoning_tokens:
                             # Force close think
                             if byte_buf:
                                 decoded = byte_buf.decode('utf-8', errors='replace')
                                 generated += decoded
                                 yield decoded
                                 byte_buf = b""
-                            # Yield closing tag
                             yield "</think>"
                             in_reasoning = False
                             think_closed = True
                             awaiting_response = True
                             response_token_count = 0
-                            # Feed the closing token
-                            ct = torch.tensor([[think_close_id]], dtype=torch.long).to(self.device)
+                            ct = torch.tensor([[think_close_id]], dtype=torch.long, device=self.device)
                             _, hidden = self.model(ct, hidden)
                             input_token = ct
                             recent_tokens.append(think_close_id)
@@ -418,7 +349,7 @@ class LinnyModel:
                     elif not in_reasoning and not think_closed:
                         response_token_count += 1
 
-                    # Send token
+                    # Output token
                     tok_str = self.tokenizer.id_to_token(nxt) or ""
                     byte_buf += _tok_to_bytes(tok_str)
                     try:
@@ -428,36 +359,31 @@ class LinnyModel:
                         byte_buf = b""
                     except UnicodeDecodeError:
                         pass
-                    
-                    input_token = torch.tensor([[nxt]], dtype=torch.long).to(self.device)
-                
+
+                    input_token = torch.tensor([[nxt]], dtype=torch.long, device=self.device)
+
                 if byte_buf:
                     leftover = byte_buf.decode('utf-8', errors='replace')
                     generated += leftover
                     yield leftover
-                
-                # Signal if we hit the token limit (for UI "keep generating" button)
-                if token_count >= max_len - 5:
-                    yield "__HIT_LIMIT__"
 
-            # ── Char generation (UPDATED with proper penalty) ──
+            # ------------------------------------------------------------------
+            # Character generation (exactly as local)
+            # ------------------------------------------------------------------
             else:
-                token_count = 0
                 for step in range(max_len):
                     logits, hidden = self.model(input_token, hidden)
                     lf = logits[0, -1].float() / max(temperature, 1e-8)
-                    
-                    # FIXED: Proper repetition penalty for char mode
+
+                    # Repetition penalty on characters
                     if penalty != 1.0 and len(recent_tokens) > 0:
-                        # Convert recent chars to indices
-                        for ch in set(recent_tokens):
-                            if isinstance(ch, str):
-                                idx = self.stoi.get(ch, 0)
-                            else:
-                                idx = ch
+                        # recent_tokens stores characters (strings)
+                        penalized_chars = set(recent_tokens)
+                        for ch in penalized_chars:
+                            idx = self.stoi.get(ch, 0)
                             if idx < lf.size(0):
                                 lf[idx] /= penalty
-                    
+
                     if top_k > 0:
                         tv, _ = torch.topk(lf, min(top_k, lf.size(-1)))
                         lf[lf < tv[-1]] = float("-inf")
@@ -468,28 +394,19 @@ class LinnyModel:
                         rm[..., 1:] = rm[..., :-1].clone()
                         rm[..., 0] = False
                         lf[si[rm]] = float("-inf")
-                    
+
                     idx = torch.multinomial(torch.softmax(lf, dim=-1), 1).item()
                     char = self.itos[idx]
-                    token_count += 1
-                    
-                    # Update recent tokens (store char for penalty)
                     recent_tokens.append(char)
-                    
-                    input_token = torch.tensor([[idx]], dtype=torch.long).to(self.device)
+                    input_token = torch.tensor([[idx]], dtype=torch.long, device=self.device)
+
                     if char == "#" and generated.endswith("##"):
                         break
-                    
                     generated += char
                     yield char
-                
-                # Signal hit limit for char mode too
-                if token_count >= max_len - 5:
-                    yield "__HIT_LIMIT__"
-
 
 # ───────────────��─────────────────────────
-# 🧰 Helpers (unchanged)
+# 🧰 Helpers (unchanged from original HF)
 # ─────────────────────────────────────────
 def parse_think_tags(text: str):
     if "<think>" not in text:
@@ -500,9 +417,7 @@ def parse_think_tags(text: str):
         return inner.strip(), (before + after).strip(), True
     return rest.strip(), before.strip(), False
 
-
 def extract_current_topic(thinking_text: str) -> str:
-    """Grab the most recent **Topic** marker from thinking text."""
     if not thinking_text:
         return "Reasoning..."
     matches = re.findall(r'\*\*([^*]+)\*\*', thinking_text)
@@ -510,51 +425,41 @@ def extract_current_topic(thinking_text: str) -> str:
         return f"Reasoning: {matches[-1].strip()}"
     return "Reasoning..."
 
-
 def format_message(visible: str, thinking: str | None,
                    think_complete: bool, think_elapsed: float | None,
                    current_topic: str = "Reasoning...") -> str:
     if not thinking:
         return visible
-
     if think_complete and think_elapsed is not None:
-        summary   = f"💭 Thought for {think_elapsed:.1f}s"
-        open_attr = ""          # collapsed when done
+        summary = f"💭 Thought for {think_elapsed:.1f}s"
+        open_attr = ""
     else:
-        summary   = current_topic
-        open_attr = ""          # always collapsed — user opens manually
-
-    think_block = (
-        f"<details class='think-details'{open_attr}>"
-        f"<summary class='think-summary'>{summary}</summary>"
-        f"<div class='think-content'>{thinking}</div>"
-        f"</details>"
-    )
+        summary = current_topic
+        open_attr = ""
+    think_block = (f"<details class='think-details'{open_attr}>"
+                   f"<summary class='think-summary'>{summary}</summary>"
+                   f"<div class='think-content'>{thinking}</div>"
+                   f"</details>")
     if visible.strip():
         return think_block + "\n\n" + visible
     return think_block
 
-
 def extract_html_canvas(text: str):
     pattern = r"```html\s*\n([\s\S]*?)```"
-    match   = re.search(pattern, text, re.IGNORECASE)
+    match = re.search(pattern, text, re.IGNORECASE)
     if match:
-        html_code    = match.group(1)
+        html_code = match.group(1)
         cleaned_text = text[:match.start()] + text[match.end():]
         return html_code.strip(), cleaned_text.strip()
     return None, text
 
-
 def make_canvas_html(code: str) -> str:
     escaped = code.replace('"', "&quot;")
-    return (
-        "<div class='canvas-wrapper'>"
-        "<div class='canvas-label'>🖼️ HTML Canvas</div>"
-        f'<iframe class="canvas-frame" srcdoc="{escaped}" '
-        'sandbox="allow-scripts" scrolling="auto"></iframe>'
-        "</div>"
-    )
-
+    return (f"<div class='canvas-wrapper'>"
+            f"<div class='canvas-label'>🖼️ HTML Canvas</div>"
+            f'<iframe class="canvas-frame" srcdoc="{escaped}" '
+            f'sandbox="allow-scripts" scrolling="auto"></iframe>'
+            f"</div>")
 
 # ─────────────────────────────────────────
 # 🚀 Auto-load
@@ -565,24 +470,18 @@ _startup_msg   = "⚠️ No model found. Place your `.pt` in the Space root."
 
 if MODEL_PATH.exists():
     try:
-        _startup_cfg   = load_config()
-        _startup_model = LinnyModel(MODEL_PATH, _startup_cfg,
-                                    tokenizer_path=TOKENIZER_PATH)
-        wm    = _startup_cfg.get("working_memory", 0)
+        _startup_cfg = load_config()
+        _startup_model = LinnyModel(MODEL_PATH, _startup_cfg, tokenizer_path=TOKENIZER_PATH)
+        wm = _startup_cfg.get("working_memory", 0)
         mtype = _startup_cfg.get("model_type", "char")
         epoch = _startup_model.epoch
-        _startup_msg = (
-            f"✅ Model auto-loaded! "
-            f"({_startup_cfg['hidden_layers']}L × {_startup_cfg['neurons']}N, "
-            f"epoch {epoch}, {mtype}"
-            + (f", {wm} ctx)" if wm else ")")
-        )
+        _startup_msg = (f"✅ Model auto-loaded! ({_startup_cfg['hidden_layers']}L × {_startup_cfg['neurons']}N, "
+                        f"epoch {epoch}, {mtype}" + (f", {wm} ctx)" if wm else ")"))
         print(_startup_msg)
     except Exception as e:
         _startup_msg = f"❌ Auto-load failed: {e}"
         print(_startup_msg)
 
-
 # ─────────────────────────────────────────
 # 🎨 CSS (unchanged)
 # ─────────────────────────────────────────
@@ -610,15 +509,12 @@ details[open] .think-summary::before { content: "▼ "; }
 .tab-nav { background: #0f0f1a !important; border-bottom: 1px solid #2d2d4e !important; }
 """
 
-
 # ─────────────────────────────────────────
-# 🖥️ UI (UPDATED with penalty window slider)
+# 🖥️ UI (identical to original HF but with penalty window)
 # ─────────────────────────────────────────
 def build_ui():
     cfg = load_config()
-
     with gr.Blocks(title="Linny AI", css=CSS) as demo:
-
         session_model = gr.State(_startup_model)
         session_cfg   = gr.State(_startup_cfg)
 
@@ -630,168 +526,97 @@ def build_ui():
         """)
 
         with gr.Tabs():
-
-            # ══════════════════════════════════════════
-            # TAB 1 — Chat
-            # ══════════════════════════════════════════
             with gr.TabItem("💬 Chat"):
-
-                model_status = gr.Markdown(
-                    value=_startup_msg,
-                    elem_classes=["status-bar"],
-                )
-
-                chatbot = gr.Chatbot(
-                    elem_id="chatbox",
-                    label="",
-                    render_markdown=True,
-                )
-
+                model_status = gr.Markdown(value=_startup_msg, elem_classes=["status-bar"])
+                chatbot = gr.Chatbot(elem_id="chatbox", label="", render_markdown=True)
                 canvas_display = gr.HTML(visible=False)
 
                 with gr.Row(elem_classes=["input-row"]):
-                    msg_box  = gr.Textbox(
-                        placeholder="Message Linny…",
-                        show_label=False,
-                        scale=8,
-                        lines=1,
-                        autofocus=True,
-                    )
+                    msg_box = gr.Textbox(placeholder="Message Linny…", show_label=False, scale=8, lines=1, autofocus=True)
                     send_btn = gr.Button("Send ↩", variant="primary", scale=1)
 
                 with gr.Row(elem_classes=["reasoning-row"]):
-                    reasoning_toggle = gr.Checkbox(
-                        value=False,
-                        label="🧠 Force Reasoning  (pre-fills <think> token)",
-                        scale=1,
-                    )
+                    reasoning_toggle = gr.Checkbox(value=False, label="🧠 Force Reasoning (pre-fills <think> token)", scale=1)
 
                 with gr.Accordion("⚙️ Generation Settings", open=False):
                     with gr.Row():
-                        temp_sl    = gr.Slider(0.1, 1.5, value=cfg.get("default_temp",    0.7),  step=0.05, label="Temperature")
-                        penalty_sl = gr.Slider(1.0, 2.0, value=cfg.get("default_penalty", 1.10), step=0.05, label="Repetition Penalty")
+                        temp_sl = gr.Slider(0.1, 1.5, value=cfg.get("default_temp", 0.8), step=0.05, label="Temperature")
+                        penalty_sl = gr.Slider(1.0, 2.0, value=cfg.get("default_penalty", 1.05), step=0.05, label="Repetition Penalty")
                     with gr.Row():
-                        penalty_window_sl = gr.Slider(10, 500, value=cfg.get("default_penalty_window", 110), step=10, label="Penalty Window (tokens)")  # ADDED
+                        penalty_window_sl = gr.Slider(10, 500, value=cfg.get("default_penalty_window", 110), step=10, label="Penalty Window")
                     with gr.Row():
-                        topp_sl    = gr.Slider(0.0, 1.0, value=cfg.get("default_top_p",   0.4),  step=0.05, label="Top-P")
-                        topk_sl    = gr.Slider(0,   100, value=cfg.get("default_top_k",   65),   step=1,    label="Top-K (0 = off)")
-                    max_len_sl = gr.Slider(100, 8000, value=cfg.get("default_max_len", 1575), step=50,
-                                           label="Max Response Length (auto-capped by Working Memory)")
+                        topp_sl = gr.Slider(0.0, 1.0, value=cfg.get("default_top_p", 0.4), step=0.05, label="Top-P")
+                        topk_sl = gr.Slider(0, 100, value=cfg.get("default_top_k", 65), step=1, label="Top-K (0 = off)")
+                    max_len_sl = gr.Slider(100, 8000, value=cfg.get("default_max_len", 2600), step=50, label="Max Response Length")
 
                 with gr.Accordion("📤 Upload Your Own Model (optional)", open=False):
                     gr.Markdown("Upload a `.pt` file and optionally a tokenizer `.json` for token-based models.")
                     with gr.Row():
-                        user_pt  = gr.File(label="Upload .pt file",              file_types=[".pt"])
+                        user_pt = gr.File(label="Upload .pt file", file_types=[".pt"])
                         user_tok = gr.File(label="Upload tokenizer .json (token model only)", file_types=[".json"])
-
-                    model_type_radio = gr.Radio(
-                        choices=["char", "token"], value="char",
-                        label="Model Type",
-                    )
-
+                    model_type_radio = gr.Radio(choices=["char", "token"], value="char", label="Model Type")
                     with gr.Row():
-                        u_layers  = gr.Number(value=5,   label="Hidden Layers", precision=0)
-                        u_neurons = gr.Number(value=768, label="Neurons",       precision=0)
-                        u_embed   = gr.Number(value=384, label="Embed Size",    precision=0)
-                        u_dropout = gr.Number(value=0.2, label="Dropout",       precision=2)
+                        u_layers = gr.Number(value=5, label="Hidden Layers", precision=0)
+                        u_neurons = gr.Number(value=768, label="Neurons", precision=0)
+                        u_embed = gr.Number(value=384, label="Embed Size", precision=0)
+                        u_dropout = gr.Number(value=0.2, label="Dropout", precision=2)
                     u_working_mem = gr.Number(value=2048, label="Working Memory (0=off)", precision=0)
                     with gr.Row():
                         u_user_tag = gr.Textbox(value="### Instruction:", label="User Tag")
-                        u_bot_tag  = gr.Textbox(value="### Response:",    label="Bot Tag")
-                        u_eos      = gr.Textbox(value="<|end|>",          label="EOS Token")
-
-                    load_user_btn    = gr.Button("🚀 Load My Model", variant="secondary")
+                        u_bot_tag = gr.Textbox(value="### Response:", label="Bot Tag")
+                        u_eos = gr.Textbox(value="<|end|>", label="EOS Token")
+                    load_user_btn = gr.Button("🚀 Load My Model", variant="secondary")
                     user_load_status = gr.Markdown("")
 
-            # ══════════════════════════════════════════
-            # TAB 2 — Admin (UPDATED with new settings)
-            # ══════════════════════════════════════════
             with gr.TabItem("🔒 Admin"):
-                gr.Markdown(
-                    "## Admin Panel\n"
-                    "Edit `config.json` settings. Upload `default_model.pt` via HF Files tab."
-                )
-
+                gr.Markdown("## Admin Panel\nEdit `config.json` settings.")
                 with gr.Row():
-                    admin_pw        = gr.Textbox(label="Admin Password", type="password",
-                                                 placeholder="Enter password", scale=3)
+                    admin_pw = gr.Textbox(label="Admin Password", type="password", placeholder="Enter password", scale=3)
                     admin_login_btn = gr.Button("Login", scale=1)
                 admin_status = gr.Markdown("")
-
                 with gr.Group(visible=False) as admin_panel:
                     gr.Markdown("### Model Architecture")
-                    model_type_admin = gr.Radio(
-                        choices=["char", "token"],
-                        value=cfg.get("model_type", "char"),
-                        label="Model Type",
-                    )
+                    model_type_admin = gr.Radio(choices=["char", "token"], value=cfg.get("model_type", "char"), label="Model Type")
                     with gr.Row():
-                        a_layers  = gr.Number(value=cfg.get("hidden_layers", 5),  label="Hidden Layers",  precision=0)
-                        a_neurons = gr.Number(value=cfg.get("neurons", 768),       label="Neurons",        precision=0)
-                        a_embed   = gr.Number(value=cfg.get("embed_size", 384),    label="Embed Size",     precision=0)
-                        a_dropout = gr.Number(value=cfg.get("dropout", 0.2),       label="Dropout",        precision=2)
-                    a_working_mem = gr.Number(value=cfg.get("working_memory", 2048),
-                                              label="Working Memory (0=off)", precision=0)
+                        a_layers = gr.Number(value=cfg.get("hidden_layers", 5), label="Hidden Layers", precision=0)
+                        a_neurons = gr.Number(value=cfg.get("neurons", 768), label="Neurons", precision=0)
+                        a_embed = gr.Number(value=cfg.get("embed_size", 384), label="Embed Size", precision=0)
+                        a_dropout = gr.Number(value=cfg.get("dropout", 0.2), label="Dropout", precision=2)
+                    a_working_mem = gr.Number(value=cfg.get("working_memory", 2048), label="Working Memory (0=off)", precision=0)
 
                     gr.Markdown("### Tags & Tokens")
                     with gr.Row():
-                        a_user_tag = gr.Textbox(value=cfg.get("user_tag",  "### Instruction:"), label="User Tag")
-                        a_bot_tag  = gr.Textbox(value=cfg.get("bot_tag",   "### Response:"),    label="Bot Tag")
-                        a_eos      = gr.Textbox(value=cfg.get("eos_token", "<|end|>"),           label="EOS Token")
-                    a_sys = gr.Textbox(
-                        value=cfg.get("system_prompt", "You are a helpful AI named Linny."),
-                        label="System Prompt", lines=2,
-                    )
+                        a_user_tag = gr.Textbox(value=cfg.get("user_tag", "### Instruction:"), label="User Tag")
+                        a_bot_tag = gr.Textbox(value=cfg.get("bot_tag", "### Response:"), label="Bot Tag")
+                        a_eos = gr.Textbox(value=cfg.get("eos_token", "<|end|>"), label="EOS Token")
+                    a_sys = gr.Textbox(value=cfg.get("system_prompt", "You are a helpful AI named Linny."), label="System Prompt", lines=2)
 
                     gr.Markdown("### Reasoning")
-                    a_reasoning_mode = gr.Radio(
-                        choices=["prompt_suffix", "response_prefix"],
-                        value=cfg.get("reasoning_mode", "prompt_suffix"),
-                        label="Force Reasoning Mode",
-                        info="prompt_suffix = appends /think to prompt  |  response_prefix = prefills <think> token",
-                    )
-                    a_reasoning_start = gr.Checkbox(
-                        value=cfg.get("reasoning_start", True),
-                        label="Enable Reasoning Start Prefix",
-                        info="Pre-fills 'I need to think about this...' after <think>",
-                    )
-                    # ADDED: New reasoning limits
-                    a_min_response_tokens = gr.Number(
-                        value=cfg.get("min_response_tokens", 3),
-                        label="Minimum Response Tokens After </think>",
-                        precision=0
-                    )
-                    a_max_reasoning_tokens = gr.Number(
-                        value=cfg.get("max_reasoning_tokens", 2500),
-                        label="Max Reasoning Tokens (auto-close if exceeded)",
-                        precision=0
-                    )
+                    a_reasoning_mode = gr.Radio(choices=["prompt_suffix", "response_prefix"], value=cfg.get("reasoning_mode", "response_prefix"), label="Force Reasoning Mode")
+                    a_reasoning_start = gr.Checkbox(value=cfg.get("reasoning_start", False), label="Enable Reasoning Start Prefix")
+                    a_min_response_tokens = gr.Number(value=cfg.get("min_response_tokens", 3), label="Minimum Response Tokens After </think>", precision=0)
+                    a_max_reasoning_tokens = gr.Number(value=cfg.get("max_reasoning_tokens", 2500), label="Max Reasoning Tokens", precision=0)
 
                     gr.Markdown("### Default Generation Settings")
                     with gr.Row():
-                        a_temp    = gr.Slider(0.1, 1.5, value=cfg.get("default_temp",    0.7),  step=0.05, label="Default Temperature")
-                        a_penalty = gr.Slider(1.0, 2.0, value=cfg.get("default_penalty", 1.10), step=0.05, label="Default Penalty")
+                        a_temp = gr.Slider(0.1, 1.5, value=cfg.get("default_temp", 0.8), step=0.05, label="Default Temperature")
+                        a_penalty = gr.Slider(1.0, 2.0, value=cfg.get("default_penalty", 1.05), step=0.05, label="Default Penalty")
                     with gr.Row():
                         a_penalty_window = gr.Slider(10, 500, value=cfg.get("default_penalty_window", 110), step=10, label="Default Penalty Window")
                     with gr.Row():
-                        a_topp    = gr.Slider(0.0, 1.0, value=cfg.get("default_top_p",   0.4),  step=0.05, label="Default Top-P")
-                        a_topk    = gr.Slider(0,   100, value=cfg.get("default_top_k",   65),   step=1,    label="Default Top-K")
-                    a_maxlen = gr.Slider(100, 8000, value=cfg.get("default_max_len", 1575), step=50, label="Default Max Length")
+                        a_topp = gr.Slider(0.0, 1.0, value=cfg.get("default_top_p", 0.4), step=0.05, label="Default Top-P")
+                        a_topk = gr.Slider(0, 100, value=cfg.get("default_top_k", 65), step=1, label="Default Top-K")
+                    a_maxlen = gr.Slider(100, 8000, value=cfg.get("default_max_len", 2600), step=50, label="Default Max Length")
 
                     save_cfg_btn = gr.Button("💾 Save config.json", variant="primary")
-                    save_status  = gr.Markdown("")
-
-        # ════════════════════════════════════════════════
-        # Callbacks (UPDATED with new fields)
-        # ════════════════════════════════════════════════
+                    save_status = gr.Markdown("")
 
+        # Callbacks
         def do_admin_login(pw):
             if pw == ADMIN_PASSWORD:
                 return gr.update(visible=True), "✅ Logged in."
             return gr.update(visible=False), "❌ Incorrect password."
-
-        admin_login_btn.click(do_admin_login, inputs=[admin_pw],
-                              outputs=[admin_panel, admin_status])
+        admin_login_btn.click(do_admin_login, inputs=[admin_pw], outputs=[admin_panel, admin_status])
 
         def do_save_config(mtype, layers, neurons, embed, dropout, working_mem,
                            user_tag, bot_tag, eos, sys_prompt,
@@ -800,94 +625,83 @@ def build_ui():
                            temp, penalty, penalty_window, top_p, top_k, max_len):
             try:
                 new_cfg = {
-                    "model_type":      mtype,
-                    "hidden_layers":   int(layers),
-                    "neurons":         int(neurons),
-                    "embed_size":      int(embed),
-                    "dropout":         float(dropout),
-                    "working_memory":  int(working_mem),
-                    "user_tag":        user_tag,
-                    "bot_tag":         bot_tag,
-                    "eos_token":       eos,
-                    "system_prompt":   sys_prompt,
-                    "reasoning_mode":  reasoning_mode,
+                    "model_type": mtype,
+                    "hidden_layers": int(layers),
+                    "neurons": int(neurons),
+                    "embed_size": int(embed),
+                    "dropout": float(dropout),
+                    "working_memory": int(working_mem),
+                    "user_tag": user_tag,
+                    "bot_tag": bot_tag,
+                    "eos_token": eos,
+                    "system_prompt": sys_prompt,
+                    "reasoning_mode": reasoning_mode,
                     "reasoning_start": bool(reasoning_start),
                     "min_response_tokens": int(min_response_tokens),
                     "max_reasoning_tokens": int(max_reasoning_tokens),
-                    "default_temp":    float(temp),
+                    "default_temp": float(temp),
                     "default_penalty": float(penalty),
                     "default_penalty_window": int(penalty_window),
-                    "default_top_p":   float(top_p),
-                    "default_top_k":   int(top_k),
+                    "default_top_p": float(top_p),
+                    "default_top_k": int(top_k),
                     "default_max_len": int(max_len),
                 }
                 save_config(new_cfg)
-                return "✅ config.json saved! Restart the Space to apply architecture changes."
+                return "✅ config.json saved! Restart the Space to apply changes."
             except Exception as e:
                 return f"❌ Error: {e}"
-
-        save_cfg_btn.click(
-            do_save_config,
+        save_cfg_btn.click(do_save_config,
             inputs=[model_type_admin, a_layers, a_neurons, a_embed, a_dropout, a_working_mem,
                     a_user_tag, a_bot_tag, a_eos, a_sys,
                     a_reasoning_mode, a_reasoning_start,
                     a_min_response_tokens, a_max_reasoning_tokens,
                     a_temp, a_penalty, a_penalty_window, a_topp, a_topk, a_maxlen],
-            outputs=[save_status],
-        )
+            outputs=[save_status])
 
         def load_user_model(pt_file, tok_file, mtype, layers, neurons, embed, dropout,
-                             working_mem, user_tag, bot_tag, eos):
+                            working_mem, user_tag, bot_tag, eos):
             if pt_file is None:
                 return None, None, "❌ Please upload a .pt file first."
             try:
                 user_cfg = {
-                    "model_type":     mtype,
-                    "hidden_layers":  int(layers),
-                    "neurons":        int(neurons),
-                    "embed_size":     int(embed),
-                    "dropout":        float(dropout),
+                    "model_type": mtype,
+                    "hidden_layers": int(layers),
+                    "neurons": int(neurons),
+                    "embed_size": int(embed),
+                    "dropout": float(dropout),
                     "working_memory": int(working_mem),
-                    "user_tag":       user_tag,
-                    "bot_tag":        bot_tag,
-                    "eos_token":      eos,
-                    "reasoning_mode":  cfg.get("reasoning_mode", "prompt_suffix"),
-                    "reasoning_start": cfg.get("reasoning_start", True),
+                    "user_tag": user_tag,
+                    "bot_tag": bot_tag,
+                    "eos_token": eos,
+                    "reasoning_mode": cfg.get("reasoning_mode", "response_prefix"),
+                    "reasoning_start": cfg.get("reasoning_start", False),
                     "min_response_tokens": cfg.get("min_response_tokens", 3),
                     "max_reasoning_tokens": cfg.get("max_reasoning_tokens", 2500),
                 }
                 tok_path = tok_file.name if tok_file else None
-                m  = LinnyModel(pt_file.name, user_cfg, tokenizer_path=tok_path)
-                wm = user_cfg["working_memory"]
-                return m, user_cfg, f"✅ Your model loaded! ({user_cfg['hidden_layers']}L × {user_cfg['neurons']}N, {mtype}, {wm} ctx, epoch {m.epoch})"
+                m = LinnyModel(pt_file.name, user_cfg, tokenizer_path=tok_path)
+                return m, user_cfg, f"✅ Model loaded! ({user_cfg['hidden_layers']}L × {user_cfg['neurons']}N, {mtype}, epoch {m.epoch})"
             except Exception as e:
                 return None, None, f"❌ Error: {e}"
-
-        load_user_btn.click(
-            load_user_model,
+        load_user_btn.click(load_user_model,
             inputs=[user_pt, user_tok, model_type_radio,
                     u_layers, u_neurons, u_embed, u_dropout, u_working_mem,
                     u_user_tag, u_bot_tag, u_eos],
-            outputs=[session_model, session_cfg, user_load_status],
-        )
+            outputs=[session_model, session_cfg, user_load_status])
 
         def respond(message, history, model, temp, penalty, penalty_window, top_p, top_k, max_len, force_thinking):
             if not message.strip():
                 yield history, "", gr.update(visible=False)
                 return
-
             if model is None:
-                yield history + [{"role":"user","content":message},{"role":"assistant","content":"⚠️ No model loaded. Place `default_model.pt` in the Space root, then restart."}], "", gr.update(visible=False)
+                yield history + [{"role":"user","content":message},{"role":"assistant","content":"⚠️ No model loaded."}], "", gr.update(visible=False)
                 return
 
             history = history + [{"role":"user","content":message},{"role":"assistant","content":""}]
-
-            eos           = model.config.get("eos_token", "<|end|>")
             full_response = ""
-            stop          = False
-            think_start   = None
-            think_end     = None
-            hit_limit     = False  # ADDED: track if we hit token limit
+            stop = False
+            think_start = None
+            think_end = None
 
             try:
                 for chunk in model.stream_generate(
@@ -895,58 +709,36 @@ def build_ui():
                     temperature=float(temp),
                     max_len=int(max_len),
                     penalty=float(penalty),
-                    penalty_window=int(penalty_window),  # ADDED
+                    penalty_window=int(penalty_window),
                     top_p=float(top_p),
                     top_k=int(top_k),
                     force_thinking=bool(force_thinking),
-                    min_response_tokens=model.config.get("min_response_tokens", 3),  # ADDED
-                    max_reasoning_tokens=model.config.get("max_reasoning_tokens", 2500),  # ADDED
+                    min_response_tokens=model.config.get("min_response_tokens", 3),
+                    max_reasoning_tokens=model.config.get("max_reasoning_tokens", 2500),
                 ):
-                    # Check for hit limit marker
-                    if chunk == "__HIT_LIMIT__":
-                        hit_limit = True
-                        continue
-                    
                     full_response += chunk
-
-                    if eos in full_response:
-                        full_response = full_response[:full_response.find(eos)]
+                    # Stop at EOS if present
+                    eos_str = model.config.get("eos_token", "<|end|>")
+                    if eos_str in full_response:
+                        full_response = full_response[:full_response.find(eos_str)]
                         if full_response.count("<think>") > full_response.count("</think>"):
                             full_response += "</think>"
                         else:
                             stop = True
 
                     thinking, visible, think_complete = parse_think_tags(full_response)
-
-                    # Timer: start on first thinking char, STOP when </think> closes
                     if thinking and think_start is None:
                         think_start = time.time()
                     if think_complete and think_end is None and think_start is not None:
                         think_end = time.time()
-
-                    if think_end is not None:
-                        think_elapsed = think_end - think_start
-                    elif think_start is not None:
-                        think_elapsed = time.time() - think_start
-                    else:
-                        think_elapsed = None
-
+                    think_elapsed = (think_end - think_start) if (think_end and think_start) else (time.time() - think_start if think_start else None)
                     current_topic = extract_current_topic(thinking) if thinking else "Reasoning..."
-
                     html_code, visible_clean = extract_html_canvas(visible)
-                    canvas_update = gr.update(
-                        visible=bool(html_code),
-                        value=make_canvas_html(html_code) if html_code else "",
-                    )
-
-                    history[-1] = {"role":"assistant","content":format_message(
-                        visible_clean, thinking, think_complete, think_elapsed, current_topic
-                    )}
+                    canvas_update = gr.update(visible=bool(html_code), value=make_canvas_html(html_code) if html_code else "")
+                    history[-1] = {"role":"assistant","content":format_message(visible_clean, thinking, think_complete, think_elapsed, current_topic)}
                     yield history, "", canvas_update
-
                     if stop:
                         break
-
             except Exception as e:
                 history[-1] = {"role":"assistant","content":f"⚠️ Generation error: {e}"}
                 yield history, "", gr.update(visible=False)
@@ -957,35 +749,21 @@ def build_ui():
             if think_end is None and think_start is not None:
                 think_end = time.time()
             think_elapsed = (think_end - think_start) if (think_end and think_start) else None
-
             if think_complete and not visible.strip():
                 visible = "*(no response generated)*"
-
-            # ADDED: Add note if hit token limit
-            if hit_limit:
-                visible += "\n\n---\n⚠️ *Response reached token limit. You can continue generating if this feature is implemented in UI.*"
-
             current_topic = extract_current_topic(thinking) if thinking else "Reasoning..."
             html_code, visible_clean = extract_html_canvas(visible)
-            history[-1] = {"role":"assistant","content":format_message(
-                visible_clean, thinking, think_complete=True,
-                think_elapsed=think_elapsed, current_topic=current_topic
-            )}
-            canvas_update = gr.update(
-                visible=bool(html_code),
-                value=make_canvas_html(html_code) if html_code else "",
-            )
+            history[-1] = {"role":"assistant","content":format_message(visible_clean, thinking, True, think_elapsed, current_topic)}
+            canvas_update = gr.update(visible=bool(html_code), value=make_canvas_html(html_code) if html_code else "")
             yield history, "", canvas_update
 
-        shared_inputs  = [msg_box, chatbot, session_model,
-                          temp_sl, penalty_sl, penalty_window_sl, topp_sl, topk_sl, max_len_sl, reasoning_toggle]
+        shared_inputs = [msg_box, chatbot, session_model,
+                         temp_sl, penalty_sl, penalty_window_sl, topp_sl, topk_sl, max_len_sl, reasoning_toggle]
         shared_outputs = [chatbot, msg_box, canvas_display]
-
-        send_btn.click(respond,  inputs=shared_inputs, outputs=shared_outputs)
+        send_btn.click(respond, inputs=shared_inputs, outputs=shared_outputs)
         msg_box.submit(respond, inputs=shared_inputs, outputs=shared_outputs)
 
     return demo
 
-
 if __name__ == "__main__":
     build_ui().launch(ssr_mode=False)