Update app.py

app.py

@@ -1,57 +1,53 @@
 # app.py – encoder-only demo for bert-beatrix-2048
-# ------------------------------------------------------------------
 # launch:  python app.py
-# ------------------------------------------------------------------
-
-import json, re, sys
+# -----------------------------------------------
+import json, re, sys, math
 from pathlib import Path, PurePosixPath
 
+import torch, torch.nn.functional as F
 import gradio as gr
 import spaces
-import torch
 from huggingface_hub import snapshot_download
 
 from bert_handler import create_handler_from_checkpoint
 
 
 # ------------------------------------------------------------------
-# 0.  Download & patch config.json  --------------------------------
-# ------------------------------------------------------------------
+# 0. Download & patch HF checkpoint --------------------------------
 REPO_ID   = "AbstractPhil/bert-beatrix-2048"
-LOCAL_DIR = "bert-beatrix-2048"
+LOCAL_CKPT = "bert-beatrix-2048"
 
-snapshot_download(REPO_ID, revision="main",
-                  local_dir=LOCAL_DIR, local_dir_use_symlinks=False)
+snapshot_download(
+    repo_id=REPO_ID,
+    revision="main",
+    local_dir=LOCAL_CKPT,
+    local_dir_use_symlinks=False,
+)
 
-cfg_path = Path(LOCAL_DIR) / "config.json"
-cfg      = json.loads(cfg_path.read_text())
+# → strip repo prefix in auto_map (one-time)
+cfg_path = Path(LOCAL_CKPT) / "config.json"
+with cfg_path.open() as f: cfg = json.load(f)
 
-auto_map, changed = cfg.get("auto_map", {}), False
-for k, v in auto_map.items():
+amap = cfg.get("auto_map", {})
+for k,v in amap.items():
     if "--" in v:
-        auto_map[k] = PurePosixPath(v.split("--", 1)[1]).as_posix()
-        changed = True
-if changed:
-    cfg["auto_map"] = auto_map
-    cfg_path.write_text(json.dumps(cfg, indent=2))
-    print("🛠️  Patched config.json → auto_map now points at local modules")
+        amap[k] = PurePosixPath(v.split("--",1)[1]).as_posix()
+cfg["auto_map"] = amap
+with cfg_path.open("w") as f: json.dump(cfg,f,indent=2)
 
-
-# ------------------------------------------------------------------
-# 1.  Model / tokenizer  -------------------------------------------
 # ------------------------------------------------------------------
-handler, full_model, tokenizer = create_handler_from_checkpoint(LOCAL_DIR)
+# 1.  Load model & components --------------------------------------
+handler, full_model, tokenizer = create_handler_from_checkpoint(LOCAL_CKPT)
 full_model = full_model.eval().cuda()
 
-encoder     = full_model.bert.encoder
-embeddings  = full_model.bert.embeddings
-emb_ln      = full_model.bert.emb_ln
-emb_drop    = full_model.bert.emb_drop
+encoder    = full_model.bert.encoder
+embeddings = full_model.bert.embeddings
+emb_ln     = full_model.bert.emb_ln
+emb_drop   = full_model.bert.emb_drop
+mlm_head   = full_model.cls          # prediction head
 
-
-# ------------------------------------------------------------------
-# 2.  Symbolic token set  ------------------------------------------
 # ------------------------------------------------------------------
+# 2. Symbolic roles -------------------------------------------------
 SYMBOLIC_ROLES = [
     "<subject>", "<subject1>", "<subject2>", "<pose>", "<emotion>",
     "<surface>", "<lighting>", "<material>", "<accessory>", "<footwear>",
@@ -60,108 +56,124 @@ SYMBOLIC_ROLES = [
     "<object_left>", "<object_right>", "<relation>", "<intent>", "<style>",
     "<fabric>", "<jewelry>",
 ]
+if any(tokenizer.convert_tokens_to_ids(t)==tokenizer.unk_token_id
+       for t in SYMBOLIC_ROLES):
+    sys.exit("❌ tokenizer missing special tokens")
 
-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}")
+# Quick helpers
+MASK = tokenizer.mask_token
 
 
 # ------------------------------------------------------------------
-# 3.  Encoder + *mask-inference* util  ------------------------------
-# ------------------------------------------------------------------
-MASK = tokenizer.mask_token or "[MASK]"
+# 3.  Encoder-plus-MLM logic ---------------------------------------
+def cosine(a,b):
+    return torch.nn.functional.cosine_similarity(a,b,dim=-1)
 
-@spaces.GPU
-def encode_and_trace(text: str, _ignored):          # all roles auto-selected
+def pool_accuracy(ids, logits, pool_mask):
     """
-    1.   run encoder pass     → cosine report (as before)
-    2.   mask **every** symbolic token one-at-a-time
-         and ask the full model to predict it back.
-         Accuracy over those positions is returned.
+    ids     : (S,)  gold token ids
+    logits  : (S,V) MLM logits
+    pool_mask : bool (S,)  which tokens belong to the candidate pool
+    returns accuracy over masked positions only (if none, return 0)
     """
-    if not text.strip():
-        return "(empty)", "0.0000", 0, "0 / 0 (0.0%)"
-
-    with torch.no_grad():
-        # --------   ENCODER PROBE (unchanged)   ------------------
-        batch = tokenizer(text, return_tensors="pt").to("cuda")
-        ids, mask = batch.input_ids, batch.attention_mask
-
-        x   = emb_drop(emb_ln(embeddings(ids)))
-        am  = full_model.bert.get_extended_attention_mask(mask, x.shape[:-1])
-        enc = encoder(x, attention_mask=am)                           # (1,S,H)
-
-        sel_ids = {tokenizer.convert_tokens_to_ids(t) for t in SYMBOLIC_ROLES}
-        flags   = torch.tensor([tid in sel_ids for tid in ids[0].tolist()],
-                               device=enc.device)
-
-        found   = [tokenizer.convert_ids_to_tokens([tid])[0]
-                   for tid in ids[0].tolist() if tid in sel_ids]
-        tokens_str = ", ".join(found) if found else "(none)"
-
-        if flags.any():
-            vec  = enc[0][flags].mean(0)
-            norm = f"{vec.norm().item():.4f}"
-        else:
-            norm = "0.0000"
-
-        # --------   MASK-AND-PREDICT ACCURACY   ------------------
-        correct, total = 0, 0
-        for pos, tid in enumerate(ids[0].tolist()):
-            if tid in sel_ids:                                        # symbolic
-                total += 1
-                masked_ids      = ids.clone()
-                masked_ids[0, pos] = tokenizer.mask_token_id
-                out = full_model(input_ids=masked_ids,
-                                 attention_mask=mask).logits          # (1,S,V)
-                pred = out[0, pos].argmax(-1).item()
-                if pred == tid:
-                    correct += 1
-
-        acc_str = f"{correct} / {total} ({(correct/total*100 if total else 0):.1f}%)"
-
-    return tokens_str, norm, len(found), acc_str
+    idx = pool_mask.nonzero(as_tuple=False).flatten()
+    if idx.numel()==0: return 0.0
+    preds = logits.argmax(-1)[idx]
+    gold  = ids[idx]
+    return (preds==gold).float().mean().item()
+
+
+@spaces.GPU
+def encode_and_trace(text, selected_roles):
+    # if user unchecked everything we treat as "all"
+    if not selected_roles:
+        selected_roles = SYMBOLIC_ROLES
+    sel_ids = {tokenizer.convert_tokens_to_ids(t) for t in selected_roles}
+
+    # ---- Tokenise & encode once ----
+    batch = tokenizer(text, return_tensors="pt").to("cuda")
+    ids, att = batch.input_ids, batch.attention_mask
+    x = emb_drop(emb_ln(embeddings(ids)))
+    ext = full_model.bert.get_extended_attention_mask(att, x.shape[:-1])
+    enc = encoder(x, attention_mask=ext)[0, :, :]      # (S,H)
+
+    # ---- compute max-cos per token (F-0/F-1) ----
+    role_mat = embeddings.word_embeddings(
+        torch.tensor(sorted(sel_ids), device=enc.device)
+    )                          # (R,H)
+    cos = cosine(enc.unsqueeze(1), role_mat.unsqueeze(0))  # (S,R)
+    maxcos, argrole = cos.max(-1)                          # (S,)
+
+    # ---- split tokens into High / Low half (F-2) ----
+    S = len(ids[0])
+    sort_idx = maxcos.argsort(descending=True)
+    hi_idx   = sort_idx[: S//2]
+    lo_idx   = sort_idx[S//2:]
+
+    # container for summary text
+    report_lines = []
+
+    # ---- pool builder helper  (uses S-4…S-7) ----
+    def greedy_pool(token_indices, direction):
+        # direction=='hi' or 'lo'
+        pool = []
+        μ = None
+        for tix in token_indices:
+            t_vec = enc[tix]
+            # incremental update (S-7)
+            μ = t_vec if μ is None else (μ*len(pool) + t_vec)/(len(pool)+1)
+            pool.append(tix)
+            # mask all tokens *except* this pool
+            masked_ids = ids.clone()
+            keep = torch.tensor(pool, device=ids.device)
+            mask_mask = torch.ones_like(ids, dtype=torch.bool)
+            mask_mask[0, keep] = False
+            masked_ids[mask_mask] = tokenizer.mask_token_id
+            # run MLM
+            with torch.no_grad():
+                logits = mlm_head(full_model.bert.emb_dl(enc.unsqueeze(0))).logits[0]
+            acc = pool_accuracy(ids[0], logits, ~mask_mask[0])
+            report_lines.append(f"{direction}-pool size {len(pool)} → acc={acc:.2f}")
+            if acc >= 0.5: break
+        return pool, acc
+
+    pool_lo, acc_lo = greedy_pool(lo_idx, "low")
+    pool_hi, acc_hi = greedy_pool(hi_idx, "high")
+
+    # ---- package textual result ----
+    res_json = {
+        "Low-pool tokens": tokenizer.decode(ids[0, pool_lo]),
+        "Low accuracy":    f"{acc_lo:.2f}",
+        "High-pool tokens":tokenizer.decode(ids[0, pool_hi]),
+        "High accuracy":   f"{acc_hi:.2f}",
+        "Trace": "\n".join(report_lines)
+    }
+    # three outputs expected by UI
+    return json.dumps(res_json, indent=2), f"{maxcos.max():.4f}", len(selected_roles)
 
 
 # ------------------------------------------------------------------
-# 4.  Gradio UI  ----------------------------------------------------
-# ------------------------------------------------------------------
+# 4.  Gradio UI -----------------------------------------------------
 def build_interface():
     with gr.Blocks(title="🧠 Symbolic Encoder Inspector") as demo:
-        gr.Markdown(
-            "## 🧠 Symbolic Encoder Inspector\n"
-            "Enter text containing the `<role>` tokens.\n"
-            "Cosine probe **and** real mask-prediction accuracy are shown."
-        )
+        gr.Markdown("## 🧠 Symbolic Encoder Inspector")
 
         with gr.Row():
             with gr.Column():
-                txt = gr.Textbox(
-                    label="Input with Symbolic Tokens",
-                    placeholder="Example: A <subject> wearing <upper_body_clothing> …",
-                    lines=3,
+                txt  = gr.Textbox(label="Prompt", lines=3)
+                roles= gr.CheckboxGroup(
+                    choices=SYMBOLIC_ROLES, label="Roles",
+                    value=SYMBOLIC_ROLES   # pre-checked
                 )
-                # checkbox group kept (pre-checked, disabled)
-                roles = gr.CheckboxGroup(
-                    choices=SYMBOLIC_ROLES,
-                    label="(all roles auto-selected)",
-                    value=SYMBOLIC_ROLES,
-                    interactive=False,
-                )
-                btn = gr.Button("Run probe + MLM check")
+                btn  = gr.Button("Run")
             with gr.Column():
-                out_tok  = gr.Textbox(label="Symbolic Tokens Found")
-                out_norm = gr.Textbox(label="Vector-norm (mean)")
-                out_cnt  = gr.Textbox(label="Token Count")
-                out_acc  = gr.Textbox(label="Mask-prediction accuracy")
-
-        btn.click(encode_and_trace,
-                  inputs=[txt, roles],
-                  outputs=[out_tok, out_norm, out_cnt, out_acc])
+                out_json = gr.Textbox(label="Result JSON")
+                out_max  = gr.Textbox(label="Max cos")
+                out_cnt  = gr.Textbox(label="# roles")
 
+        btn.click(encode_and_trace, [txt,roles], [out_json,out_max,out_cnt])
     return demo
 
 
-if __name__ == "__main__":
-    build_interface().launch()
+if __name__=="__main__":
+    build_interface().launch()