Update app.py

app.py

@@ -1,10 +1,18 @@
+# app.py
+# --------------------------------------------------------------------------------------------------
+# Gradio app for Beeper
+# - Loads released safetensors + tokenizer from Hugging Face
+# - Auto-sizes pentachora banks to match checkpoints (across Beeper v1..v4)
+# - Generation uses same knobs & penalties as training script
+# --------------------------------------------------------------------------------------------------
 import gradio as gr
 import torch
-from beeper_model import BeeperRoseGPT, generate
 from tokenizers import Tokenizer
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file as load_safetensors
 
+from beeper import BeeperRoseGPT, generate, prepare_model_for_state_dict
+
 # ----------------------------
 # 🔧 Model versions configuration
 # ----------------------------
@@ -31,8 +39,8 @@ MODEL_VERSIONS = {
     },
 }
 
-# Base configuration
-config = {
+# Base configuration (matches training defaults)
+CONFIG = {
     "context": 512,
     "vocab_size": 8192,
     "dim": 512,
@@ -42,171 +50,169 @@ config = {
     "temperature": 0.9,
     "top_k": 40,
     "top_p": 0.9,
-    "repetition_penalty": 1.1,
+    "repetition_penalty": 1.10,
     "presence_penalty": 0.6,
     "frequency_penalty": 0.0,
     "resid_dropout": 0.1,
     "dropout": 0.0,
     "grad_checkpoint": False,
-    "tokenizer_path": "beeper.tokenizer.json"
+    # tokenizer_path not needed here; we load tokenizer.json from the HF repo
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Global model and tokenizer variables
-infer = None
-tok = None
-current_version = None
+# Globals (kept simple for a single process Gradio app)
+infer: BeeperRoseGPT | None = None
+tok: Tokenizer | None = None
+current_version: str | None = None
+
 
-def load_model_version(version_name):
-    """Load the selected model version"""
+def load_model_version(version_name: str) -> str:
+    """
+    Download the checkpoint and tokenizer, build model, ensure pentachora sizes match,
+    then strictly load weights. Robust to v1/v2 (no pentas) and v3/v4 (with pentas).
+    """
     global infer, tok, current_version
-    
-    if current_version == version_name and infer is not None:
+
+    if current_version == version_name and infer is not None and tok is not None:
         return f"Already loaded: {version_name}"
-    
+
     version_info = MODEL_VERSIONS[version_name]
-    
+
     try:
-        # Download model and tokenizer files
+        # Download artifacts
         model_file = hf_hub_download(
-            repo_id=version_info["repo_id"], 
+            repo_id=version_info["repo_id"],
             filename=version_info["model_file"]
         )
         tokenizer_file = hf_hub_download(
-            repo_id=version_info["repo_id"], 
+            repo_id=version_info["repo_id"],
             filename="tokenizer.json"
         )
-        
-        # Initialize model
-        infer = BeeperRoseGPT(config).to(device)
-        
-        # Load safetensors
-        state_dict = load_safetensors(model_file, device=str(device))
-        infer.load_state_dict(state_dict)
-        infer.eval()
-        
-        # Load tokenizer
-        tok = Tokenizer.from_file(tokenizer_file)
-        
+
+        # Load state dict on CPU, inspect pentachora shapes if present
+        state_dict = load_safetensors(model_file, device="cpu")
+
+        # Build model & pre-create pentachora if needed
+        m = BeeperRoseGPT(CONFIG).to(device)
+        prepare_model_for_state_dict(m, state_dict, device=device)
+
+        # Try strict load first; if shapes drift (rare), fallback to non-strict
+        try:
+            missing, unexpected = m.load_state_dict(state_dict, strict=True)
+            # PyTorch returns NamedTuple; report counts
+            _msg = f"strict load ok | missing={len(missing)} unexpected={len(unexpected)}"
+        except Exception as e:
+            _msg = f"strict load failed ({e}); trying non-strict…"
+            # Non-strict load for very old snapshots
+            m.load_state_dict(state_dict, strict=False)
+
+        m.eval()
+
+        # Tokenizer
+        t = Tokenizer.from_file(tokenizer_file)
+
+        # Swap globals
+        infer, tok = m, t
         current_version = version_name
-        return f"Successfully loaded: {version_name}"
-    
+        return f"Successfully loaded: {version_name} ({_msg})"
+
     except Exception as e:
+        infer = None
+        tok = None
+        current_version = None
         return f"Error loading {version_name}: {str(e)}"
 
-# Load default model on startup - try v4 first, fallback to v3
+
+# Load default on startup — prefer v4, fallback to v3
 try:
     load_status = load_model_version("Beeper v4 (Advanced)")
     if "Error" in load_status:
         print(f"v4 not ready yet: {load_status}")
         load_status = load_model_version("Beeper v3 (Multi-Concept)")
-except:
+except Exception as _:
     load_status = load_model_version("Beeper v3 (Multi-Concept)")
-    
 print(load_status)
 
+
 # ----------------------------
-# ���� Gradio Chat Wrapper
+# 💬 Chat wrapper
 # ----------------------------
-def beeper_reply(message, history, model_version, temperature=None, top_k=None, top_p=None, max_new_tokens=80):
+def beeper_reply(
+    message: str,
+    history: list[tuple[str, str]] | None,
+    model_version: str,
+    temperature: float | None,
+    top_k: int | None,
+    top_p: float | None,
+    max_new_tokens: int = 80
+) -> str:
     global infer, tok, current_version
-    
-    # Load model if version changed
+
+    # Hot-swap versions if the dropdown changed
     if model_version != current_version:
         status = load_model_version(model_version)
         if "Error" in status:
             return f"⚠️ {status}"
-    
-    # Check if model is loaded
+
     if infer is None or tok is None:
         return "⚠️ Model not loaded. Please select a version and try again."
-    
-    # Use defaults if not provided
-    if temperature is None:
-        temperature = 0.9
-    if top_k is None:
-        top_k = 40
-    if top_p is None:
-        top_p = 0.9
-    
-    # Try Q&A format since she has some in corpus
-    if "?" in message:
-        prompt = f"Q: {message}\nA:"
-    elif message.lower().strip() in ["hi", "hello", "hey"]:
-        prompt = "The little robot said hello. She said, \""
-    elif "story" in message.lower():
+
+    # Light prompting heuristics (consistent with your example)
+    m = message.strip()
+    if "?" in m:
+        prompt = f"Q: {m}\nA:"
+    elif m.lower() in {"hi", "hello", "hey"}:
+        prompt = 'The little robot said hello. She said, "'
+    elif "story" in m.lower():
         prompt = "Once upon a time, there was a robot. "
     else:
-        # Simple continuation
-        prompt = message + ". "
-    
-    # Generate response with lower temperature for less repetition
-    response = generate(
+        prompt = m + ". "
+
+    # Generate
+    text = generate(
         model=infer,
         tok=tok,
-        cfg=config,
+        cfg=CONFIG,
         prompt=prompt,
-        max_new_tokens=max_new_tokens,  # Shorter to avoid rambling
-        temperature=float(temperature),  # Slightly lower temp
-        top_k=int(top_k),
-        top_p=float(top_p),
-        repetition_penalty=1.1,  # Higher penalty for repetition
-        presence_penalty=0.8,    # Higher presence penalty
-        frequency_penalty=0.1,    # Add frequency penalty
+        max_new_tokens=int(max_new_tokens),
+        temperature=float(temperature) if temperature is not None else None,
+        top_k=int(top_k) if top_k is not None else None,
+        top_p=float(top_p) if top_p is not None else None,
+        repetition_penalty=1.10,
+        presence_penalty=0.8,
+        frequency_penalty=0.1,
         device=device,
-        detokenize=True
+        detokenize=True,
     )
-    
-    # Aggressive cleanup
-    # Remove the prompt completely
-    if response.startswith(prompt):
-        response = response[len(prompt):]
-    
-    # Remove Q&A format artifacts
-    response = response.replace("Q:", "").replace("A:", "")
-    
-    # Split on newlines and take first non-empty line
-    lines = response.split('\n')
-    for line in lines:
-        clean_line = line.strip()
-        if clean_line and not clean_line.startswith(message[:10]):
-            response = clean_line
-            break
-    
-    # If response still contains the user message, try to extract after it
-    if message.lower()[:20] in response.lower()[:50]:
-        # Find where the echo ends
-        words_in_message = message.split()
-        for i in range(min(5, len(words_in_message)), 0, -1):
-            pattern = ' '.join(words_in_message[:i])
-            if pattern.lower() in response.lower():
-                idx = response.lower().find(pattern.lower()) + len(pattern)
-                response = response[idx:].strip()
-                break
-    
-    # Remove any remaining "User" or "Beeper" artifacts
-    for artifact in ["User:", "Beeper:", "U ser:", "Beep er:", "User ", "Beeper "]:
-        response = response.replace(artifact, "")
-    
-    # Ensure we have something
-    if not response or len(response) < 3:
-        responses = [
-            "I like robots and stories!",
-            "That's interesting!",
-            "I want to play in the park.",
-            "The robot was happy.",
-            "Yes, I think so too!"
-        ]
-        import random
-        response = random.choice(responses)
-    
-    # Clean ending
-    response = response.strip()
-    if response and response[-1] not in '.!?"':
-        response = response.rsplit('.', 1)[0] + '.' if '.' in response else response + '.'
-    
-    return response[:200]  # Cap length
+
+    # Strip prompt echoes & artifacts
+    if text.startswith(prompt):
+        text = text[len(prompt):]
+    text = text.replace("Q:", "").replace("A:", "")
+
+    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
+    if lines:
+        text = lines[0]
+
+    # If user message echoed at head, trim after first occurrence
+    head = m[:20].lower()
+    if text.lower().startswith(head):
+        idx = text.lower().find(head)
+        text = text[idx + len(head):].strip() or text
+
+    for artifact in ("User:", "Beeper:", "U ser:", "Beep er:", "User ", "Beeper "):
+        text = text.replace(artifact, "")
+
+    text = text.strip()
+    if not text or len(text) < 3:
+        text = "I like robots and stories!"
+
+    if text[-1:] not in ".!?”\"'":
+        text += "."
+
+    return text[:200]
+
 
 # ----------------------------
 # 🖼️ Interface
@@ -214,37 +220,34 @@ def beeper_reply(message, history, model_version, temperature=None, top_k=None,
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
-        # 🤖 Beeper - A Rose-based Tiny Language Model
-        Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me - I'm still learning! 💕
+        # 🤖 Beeper — A Rose-based Tiny Language Model
+        Hello! I'm Beeper, a small language model trained with love and care. Please be patient with me — I'm still learning! 💕
         """
     )
-    
+
     with gr.Row():
         with gr.Column(scale=3):
             model_dropdown = gr.Dropdown(
                 choices=list(MODEL_VERSIONS.keys()),
-                value="Beeper v3 (Multi-Concept)",  # Default to v3 since v4 might not be ready
+                value="Beeper v3 (Multi-Concept)",  # safer default
                 label="Select Beeper Version",
-                info="Choose which version of Beeper to chat with"
+                info="Choose which version of Beeper to chat with",
             )
         with gr.Column(scale=7):
-            version_info = gr.Markdown("**Current:** Beeper v3 with 30+ epochs including reasoning, math, and ethics")
-    
-    # Update version info when dropdown changes
-    def update_version_info(version_name):
-        info = MODEL_VERSIONS[version_name]["description"]
-        return f"**Current:** {info}"
-    
+            version_info = gr.Markdown("**Current:** " + MODEL_VERSIONS["Beeper v3 (Multi-Concept)"]["description"])
+
+    def update_version_info(version_name: str):
+        return f"**Current:** {MODEL_VERSIONS[version_name]['description']}"
+
     model_dropdown.change(
         fn=update_version_info,
         inputs=[model_dropdown],
-        outputs=[version_info]
+        outputs=[version_info],
     )
-    
-    # Chat interface
-    chatbot = gr.Chatbot(label="Chat with Beeper", type="tuples", height=400)
+
+    chatbot = gr.Chatbot(label="Chat with Beeper", height=400)
     msg = gr.Textbox(label="Message", placeholder="Type your message here...")
-    
+
     with gr.Row():
         with gr.Column(scale=2):
             temperature_slider = gr.Slider(0.1, 1.5, value=0.9, step=0.1, label="Temperature")
@@ -253,13 +256,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=2):
             top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p")
         with gr.Column(scale=2):
-            max_new_tokens_slider = gr.Slider(20, 512, value=128, step=1, label="Max-new-tokens")
-    
+            max_new_tokens_slider = gr.Slider(20, 512, value=128, step=1, label="Max new tokens")
+
     with gr.Row():
         submit = gr.Button("Send", variant="primary")
         clear = gr.Button("Clear")
-    
-    # Examples
+
     gr.Examples(
         examples=[
             ["Hello Beeper! How are you today?"],
@@ -268,28 +270,27 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             ["What makes you happy?"],
             ["Tell me about your dreams"],
         ],
-        inputs=msg
+        inputs=msg,
     )
-    
-    # Handle chat
+
     def respond(message, chat_history, model_version, temperature, top_k, top_p, max_new_tokens):
-        if not chat_history:
+        if chat_history is None:
             chat_history = []
         response = beeper_reply(message, chat_history, model_version, temperature, top_k, top_p, max_new_tokens)
-        chat_history.append([message, response])
+        chat_history.append((message, response))
         return "", chat_history
-    
+
     msg.submit(
-        respond, 
-        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider, max_new_tokens_slider], 
-        [msg, chatbot]
+        respond,
+        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider, max_new_tokens_slider],
+        [msg, chatbot],
     )
     submit.click(
-        respond, 
-        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider, max_new_tokens_slider], 
-        [msg, chatbot]
+        respond,
+        [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider, max_new_tokens_slider],
+        [msg, chatbot],
     )
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()