Update app.py

app.py

@@ -1,10 +1,7 @@
 # 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
-# --------------------------------------------------------------------------------------------------
+# Gradio app exposing full Corpus (coarse) and Capoera (topic/mood) selections
+
+import json
 import gradio as gr
 import torch
 from tokenizers import Tokenizer
@@ -13,9 +10,6 @@ from safetensors.torch import load_file as load_safetensors
 
 from beeper_model import BeeperRoseGPT, generate, prepare_model_for_state_dict
 
-# ----------------------------
-# 🔧 Model versions configuration
-# ----------------------------
 MODEL_VERSIONS = {
     "Beeper v4 (Advanced)": {
         "repo_id": "AbstractPhil/beeper-rose-v4",
@@ -39,7 +33,6 @@ MODEL_VERSIONS = {
     },
 }
 
-# Base configuration (matches training defaults)
 CONFIG = {
     "context": 512,
     "vocab_size": 8192,
@@ -56,196 +49,182 @@ CONFIG = {
     "resid_dropout": 0.1,
     "dropout": 0.0,
     "grad_checkpoint": False,
-    # tokenizer_path not needed here; we load tokenizer.json from the HF repo
+    "runtime_pentachora": {
+        "enable": True,
+        "pool": "mean",
+        "temp": 0.10,
+        "coarse_alpha": 0.25,
+        "topic_alpha":  0.15,
+        "mood_alpha":   0.10,
+    },
 }
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Globals (kept simple for a single process Gradio app)
 infer: BeeperRoseGPT | None = None
 tok: Tokenizer | None = None
 current_version: str | None = None
 
+# Metadata for selectors
+CORPUS_CHOICES: list[str] = []
+CORPUS_INDEX: dict[str, int] = {}
+TOPIC_CHOICES: list[str] = []
+MOOD_CHOICES: list[str] = []
+
+def _mood_labels(mood_bins: int) -> list[str]:
+    center = mood_bins // 2
+    labels = []
+    for i in range(mood_bins):
+        v = i - center
+        name = { -3:"Very Negative", -2:"Negative", -1:"Slightly Negative",
+                  0:"Neutral", 1:"Slightly Positive", 2:"Positive", 3:"Very Positive" }.get(v, f"Valence {v:+d}")
+        labels.append(f"{i} ({name} {v:+d})")
+    return labels
+
+def _build_choices_from_config(repo_id: str, coarse_C: int, topic_C: int, mood_C: int):
+    global CORPUS_CHOICES, CORPUS_INDEX, TOPIC_CHOICES, MOOD_CHOICES
+    CORPUS_CHOICES, CORPUS_INDEX = [], {}
+    # Try to load training config.json (exported alongside weights)
+    names = []
+    try:
+        cfg_path = hf_hub_download(repo_id, "config.json")
+        with open(cfg_path, "r", encoding="utf-8") as f:
+            train_cfg = json.load(f)
+        alive = train_cfg.get("_alive_entries")
+        if isinstance(alive, list) and all(isinstance(e, dict) for e in alive):
+            names = [str(e.get("name", f"Class {i}")) for i, e in enumerate(alive)]
+        elif isinstance(train_cfg.get("corpus"), list):
+            # fallback: use corpus list if length matches bank size
+            maybe = [str(e.get("name", f"Class {i}")) for i, e in enumerate(train_cfg["corpus"])]
+            if len(maybe) == coarse_C:
+                names = maybe
+    except Exception:
+        names = []
+
+    if len(names) != coarse_C:
+        names = [f"Class {i}" for i in range(coarse_C)]
+
+    CORPUS_CHOICES = names
+    CORPUS_INDEX = {name: i for i, name in enumerate(names)}
+    TOPIC_CHOICES = [str(i) for i in range(topic_C)]
+    MOOD_CHOICES  = _mood_labels(mood_C)
 
 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
-
+    global infer, tok, current_version, CORPUS_CHOICES, TOPIC_CHOICES, MOOD_CHOICES
     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]
-
+    info = MODEL_VERSIONS[version_name]
     try:
-        # Download artifacts
-        model_file = hf_hub_download(
-            repo_id=version_info["repo_id"],
-            filename=version_info["model_file"]
-        )
-        tokenizer_file = hf_hub_download(
-            repo_id=version_info["repo_id"],
-            filename="tokenizer.json"
-        )
-
-        # Load state dict on CPU, inspect pentachora shapes if present
-        state_dict = load_safetensors(model_file, device="cpu")
+        model_file = hf_hub_download(info["repo_id"], info["model_file"])
+        tokenizer_file = hf_hub_download(info["repo_id"], "tokenizer.json")
 
-        # Build model & pre-create pentachora if needed
+        state = load_safetensors(model_file, device="cpu")
         m = BeeperRoseGPT(CONFIG).to(device)
-        prepare_model_for_state_dict(m, state_dict, device=device)
+        prepare_model_for_state_dict(m, state, 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
+            missing, unexpected = m.load_state_dict(state, strict=True)
             _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)
+            _msg = f"strict load failed ({e}); non-strict fallback"
+            m.load_state_dict(state, 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} ({_msg})"
+        infer, tok, current_version = m, t, version_name
+
+        # Build UI choices from bank sizes + training config (for names)
+        coarse_C = infer.penta_coarse.size(0) if infer.penta_coarse is not None else 0
+        topic_C  = infer.penta_medium.size(0) if infer.penta_medium is not None else 512
+        mood_C   = infer.penta_fine.size(0) if infer.penta_fine is not None else 7
+        _build_choices_from_config(info["repo_id"], coarse_C, topic_C, mood_C)
 
+        return f"Successfully loaded: {version_name} ({_msg})"
     except Exception as e:
-        infer = None
-        tok = None
-        current_version = None
+        infer = None; tok = None; current_version = None
+        CORPUS_CHOICES, TOPIC_CHOICES, MOOD_CHOICES = [], [], []
         return f"Error loading {version_name}: {str(e)}"
 
-
-# Load default on startup — prefer v4, fallback to v3
+# Initial load: 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 Exception as _:
-    load_status = load_model_version("Beeper v3 (Multi-Concept)")
-print(load_status)
-
-
-# ----------------------------
-# 💬 Chat wrapper
-# ----------------------------
-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:
+    status = load_model_version("Beeper v4 (Advanced)")
+    if "Error" in status:
+        print(status)
+        status = load_model_version("Beeper v3 (Multi-Concept)")
+except Exception:
+    status = load_model_version("Beeper v3 (Multi-Concept)")
+print(status)
+
+def _parse_selected_indices(values: list[str] | None, mapping: dict[str,int] | None = None) -> list[int] | None:
+    if not values: return None
+    if mapping is None:
+        return [int(v.split()[0]) if isinstance(v, str) else int(v) for v in values]
+    return [mapping[v] for v in values if v in mapping]
+
+def beeper_reply(message, history, model_version, temperature, top_k, top_p, max_new_tokens,
+                 corpus_selected, topic_selected, mood_selected):
     global infer, tok, current_version
-
-    # 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}"
+        s = load_model_version(model_version)
+        if "Error" in s:
+            return f"⚠️ {s}"
 
     if infer is None or tok is None:
         return "⚠️ Model not loaded. Please select a version and try again."
 
-    # 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:
-        prompt = m + ". "
-
-    # Generate
-    text = generate(
-        model=infer,
-        tok=tok,
-        cfg=CONFIG,
-        prompt=prompt,
+    # Build runtime pull config with user selections
+    rt = dict(CONFIG.get("runtime_pentachora", {}))
+    # Convert selections -> index lists
+    rt["coarse_select"] = _parse_selected_indices(corpus_selected, CORPUS_INDEX)  # names -> indices
+    rt["topic_select"]  = _parse_selected_indices(topic_selected, None)          # numeric strings -> ints
+    rt["mood_select"]   = _parse_selected_indices(mood_selected, None)           # numeric strings -> ints
+
+    m = (message or "").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:              prompt = m + ". "
+
+    out = generate(
+        model=infer, tok=tok, cfg=CONFIG, prompt=prompt,
         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,
+        repetition_penalty=1.10, presence_penalty=0.8, frequency_penalty=0.1,
+        device=device, detokenize=True, runtime_cfg=rt,
     )
 
-    # 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, "")
+    if out.startswith(prompt): out = out[len(prompt):]
+    out = out.replace("Q:","").replace("A:","").strip()
+    if out and out[-1] not in ".!?”\"'": out += "."
+    return out[:200]
 
-    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
-# ----------------------------
+# ---------------- UI ----------------
 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! 💕
-        """
-    )
+    gr.Markdown("# 🤖 Beeper — Corpus & Capoera–aware Chat")
 
     with gr.Row():
         with gr.Column(scale=3):
             model_dropdown = gr.Dropdown(
                 choices=list(MODEL_VERSIONS.keys()),
-                value="Beeper v3 (Multi-Concept)",  # safer default
-                label="Select Beeper Version",
-                info="Choose which version of Beeper to chat with",
+                value="Beeper v3 (Multi-Concept)",
+                label="Select Beeper Version"
             )
         with gr.Column(scale=7):
             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],
-    )
-
-    chatbot = gr.Chatbot(label="Chat with Beeper", height=400)
+    # Runtime pentachora selectors
+    with gr.Row():
+        with gr.Column():
+            corpus_select = gr.Dropdown(choices=CORPUS_CHOICES, multiselect=True, label="Corpus (Coarse classes)")
+        with gr.Column():
+            topic_select = gr.Dropdown(choices=TOPIC_CHOICES, multiselect=True, label="Capoera Topics (IDs)")
+        with gr.Column():
+            mood_select  = gr.Dropdown(choices=MOOD_CHOICES,  multiselect=True, label="Capoera Moods (valence)")
+
+    chatbot = gr.Chatbot(label="Chat with Beeper", height=420)
     msg = gr.Textbox(label="Message", placeholder="Type your message here...")
 
     with gr.Row():
@@ -262,34 +241,38 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         submit = gr.Button("Send", variant="primary")
         clear = gr.Button("Clear")
 
-    gr.Examples(
-        examples=[
-            ["Hello Beeper! How are you today?"],
-            ["Can you tell me a story about a robot?"],
-            ["What do you like to do for fun?"],
-            ["What makes you happy?"],
-            ["Tell me about your dreams"],
-        ],
-        inputs=msg,
+    # On version change: load model + update selectors
+    def on_change_version(version_name: str):
+        status = load_model_version(version_name)
+        info = f"**Current:** {MODEL_VERSIONS[version_name]['description']}  \n{status}"
+        # refresh selector choices
+        return (
+            info,
+            gr.update(choices=CORPUS_CHOICES, value=[]),
+            gr.update(choices=TOPIC_CHOICES,  value=[]),
+            gr.update(choices=MOOD_CHOICES,   value=[]),
+        )
+
+    model_dropdown.change(
+        on_change_version,
+        inputs=[model_dropdown],
+        outputs=[version_info, corpus_select, topic_select, mood_select],
     )
 
-    def respond(message, chat_history, model_version, temperature, top_k, top_p, max_new_tokens):
-        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))
+    def respond(message, chat_history, model_version, temperature, top_k, top_p, max_new_tokens,
+                corpus_selected, topic_selected, mood_selected):
+        if chat_history is None: chat_history = []
+        resp = beeper_reply(message, chat_history, model_version, temperature, top_k, top_p, max_new_tokens,
+                            corpus_selected, topic_selected, mood_selected)
+        chat_history.append((message, resp))
         return "", chat_history
 
-    msg.submit(
-        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],
-    )
+    inputs_all = [msg, chatbot, model_dropdown, temperature_slider, top_k_slider, top_p_slider, max_new_tokens_slider,
+                  corpus_select, topic_select, mood_select]
+    outputs_all = [msg, chatbot]
+
+    msg.submit(respond, inputs_all, outputs_all)
+    submit.click(respond, inputs_all, outputs_all)
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":