Update app.py

app.py

@@ -1,13 +1,9 @@
 # app.py – encoder-only demo for bert-beatrix-2048
-# -----------------------------------------------
-# launch:  python app.py
-# (gradio UI appears at http://localhost:7860)
-
-import json
-import re
-import sys
-from pathlib import Path, PurePosixPath   # ← PurePosixPath import added
+# ------------------------------------------------
+# launch:  python app.py     →  http://localhost:7860
 
+import json, re, sys
+from pathlib import Path, PurePosixPath        # ← PurePosixPath import
 import gradio as gr
 import spaces
 import torch
@@ -17,10 +13,10 @@ from bert_handler import create_handler_from_checkpoint
 
 
 # ------------------------------------------------------------------
-# 0.  Download & patch config.json  --------------------------------
+# 0.  Download & patch config.json  ---------------------------------
 # ------------------------------------------------------------------
-REPO_ID = "AbstractPhil/bert-beatrix-2048"
-LOCAL_CKPT = "bert-beatrix-2048"          # cached dir name
+REPO_ID     = "AbstractPhil/bert-beatrix-2048"
+LOCAL_CKPT  = "bert-beatrix-2048"              # cache dir name
 
 snapshot_download(
     repo_id=REPO_ID,
@@ -29,32 +25,30 @@ snapshot_download(
     local_dir_use_symlinks=False,
 )
 
-# ── one-time patch: strip the “repo--” prefix that confuses AutoModel ──
 cfg_path = Path(LOCAL_CKPT) / "config.json"
 with cfg_path.open() as f:
     cfg = json.load(f)
 
 auto_map = cfg.get("auto_map", {})
-changed = False
+patched  = False
 for k, v in auto_map.items():
-    if "--" in v:                         # v looks like  "repo--module.Class"
+    if "--" in v:                              # strip   repo--module.Class
         auto_map[k] = PurePosixPath(v.split("--", 1)[1]).as_posix()
-        changed = True
+        patched = True
 
-if changed:
+if patched:
     cfg["auto_map"] = auto_map
     with cfg_path.open("w") as f:
         json.dump(cfg, f, indent=2)
-    print("🛠️  Patched config.json → auto_map now points at local modules")
+    print("🛠️  Patched config.json → auto_map now points to local modules")
 
 
 # ------------------------------------------------------------------
-# 1.  Model / tokenizer  -------------------------------------------
+# 1.  Load model / tokenizer  ---------------------------------------
 # ------------------------------------------------------------------
 handler, full_model, tokenizer = create_handler_from_checkpoint(LOCAL_CKPT)
 full_model = full_model.eval().cuda()
 
-# Grab encoder + embedding stack only
 encoder     = full_model.bert.encoder
 embeddings  = full_model.bert.embeddings
 emb_ln      = full_model.bert.emb_ln
@@ -62,7 +56,7 @@ emb_drop    = full_model.bert.emb_drop
 
 
 # ------------------------------------------------------------------
-# 2.  Symbolic token set  ------------------------------------------
+# 2.  Symbolic roles  ------------------------------------------------
 # ------------------------------------------------------------------
 SYMBOLIC_ROLES = [
     "<subject>", "<subject1>", "<subject2>", "<pose>", "<emotion>",
@@ -73,44 +67,47 @@ SYMBOLIC_ROLES = [
     "<fabric>", "<jewelry>",
 ]
 
-# quick sanity check
-missing = [tok for tok in SYMBOLIC_ROLES
-           if tokenizer.convert_tokens_to_ids(tok) == tokenizer.unk_token_id]
+missing = [t for t in SYMBOLIC_ROLES
+           if tokenizer.convert_tokens_to_ids(t) == tokenizer.unk_token_id]
 if missing:
     sys.exit(f"❌ Tokenizer is missing {missing}")
 
 
 # ------------------------------------------------------------------
-# 3.  Encoder-only inference util  ---------------------------------
+# 3.  Encoder-only helper  ------------------------------------------
 # ------------------------------------------------------------------
 @spaces.GPU
 def encode_and_trace(text: str, selected_roles: list[str]):
+    """
+    Returns **exactly three scalars** matching the 3 gradio outputs:
+        1) tokens_str  (comma-separated list found)
+        2) norm_str    (mean L2-norm of those embeddings)
+        3) count_int   (# tokens matched)
+    """
     with torch.no_grad():
         batch = tokenizer(text, return_tensors="pt").to("cuda")
-        ids, mask = batch.input_ids, batch.attention_mask
+        ids, attn = batch.input_ids, batch.attention_mask
 
         x = emb_drop(emb_ln(embeddings(ids)))
-        ext_mask = full_model.bert.get_extended_attention_mask(mask, x.shape[:-1])
-        enc = encoder(x, attention_mask=ext_mask)              # (1, S, H)
+        ext = full_model.bert.get_extended_attention_mask(attn, x.shape[:-1])
+        enc = encoder(x, attention_mask=ext)                   # (1, S, H)
 
-        sel_ids = {tokenizer.convert_tokens_to_ids(t) for t in selected_roles}
-        flags   = torch.tensor([tid in sel_ids for tid in ids[0].tolist()],
-                               device=enc.device)
+        role_ids = {tokenizer.convert_tokens_to_ids(t) for t in selected_roles}
+        mask     = torch.tensor([tid in role_ids for tid in ids[0].tolist()],
+                                device=enc.device, dtype=torch.bool)
 
-        found_tokens = [tokenizer.convert_ids_to_tokens([tid])[0]
-                        for tid in ids[0].tolist() if tid in sel_ids]
-        tokens_str = ", ".join(found_tokens) or "(none)"
+        found = [tokenizer.convert_ids_to_tokens([tid])[0]
+                 for tid in ids[0].tolist() if tid in role_ids]
+        tokens_str = ", ".join(found) or "(none)"
 
-        if flags.any():
-            vec  = enc[0][flags].mean(0)
-            norm = f"{vec.norm().item():.4f}"
+        if mask.any():
+            mean_vec = enc[0][mask].mean(0)
+            norm_str = f"{mean_vec.norm().item():.4f}"
         else:
-            norm = "0.0000"
-
-        count = int(flags.sum().item())
-        # >>> return *three* scalars, not one dict  <<<
-        return tokens_str, norm, count
+            norm_str = "0.0000"
 
+        count_int = int(mask.sum().item())
+        return tokens_str, norm_str, count_int       # ← three outputs!
 
 
 # ------------------------------------------------------------------
@@ -119,29 +116,23 @@ def encode_and_trace(text: str, selected_roles: list[str]):
 def build_interface():
     with gr.Blocks(title="🧠 Symbolic Encoder Inspector") as demo:
         gr.Markdown(
-            "## 🧠 Symbolic Encoder Inspector\n"
-            "Paste some text containing the special `<role>` tokens and "
-            "inspect their encoder representations."
+            "### 🧠 Symbolic Encoder Inspector\n"
+            "Paste text that includes the `<role>` tokens and inspect their "
+            "hidden-state statistics."
         )
 
         with gr.Row():
             with gr.Column():
-                txt = gr.Textbox(
-                    label="Input with Symbolic Tokens",
-                    placeholder="Example: A <subject> wearing <upper_body_clothing> …",
-                    lines=3,
-                )
-                roles = gr.CheckboxGroup(
-                    choices=SYMBOLIC_ROLES,
-                    label="Trace these symbolic roles",
-                )
-                btn = gr.Button("Encode & Trace")
+                txt  = gr.Textbox(label="Input", lines=3,
+                                  placeholder="A <subject> wearing <upper_body_clothing> …")
+                chk  = gr.CheckboxGroup(SYMBOLIC_ROLES, label="Roles to trace")
+                run  = gr.Button("Encode & Trace")
             with gr.Column():
-                out_tok  = gr.Textbox(label="Symbolic Tokens Found")
-                out_norm = gr.Textbox(label="Mean Norm")
-                out_cnt  = gr.Textbox(label="Token Count")
+                out_tok  = gr.Textbox(label="Tokens found")
+                out_norm = gr.Textbox(label="Mean norm")
+                out_cnt  = gr.Textbox(label="Token count")
 
-        btn.click(encode_and_trace, [txt, roles], [out_tok, out_norm, out_cnt])
+        run.click(encode_and_trace, [txt, chk], [out_tok, out_norm, out_cnt])
 
     return demo