Update app.py

app.py

@@ -113,49 +113,41 @@ def encode_and_trace(text, selected_roles):
     # container for summary text
     report_lines = []
     
-    # ------------------------------------------------------------------
-    # Greedy pool helper  – tensor-safe version
-    # ------------------------------------------------------------------
-    def greedy_pool(index_tensor: torch.Tensor, which: str):
+    # ───────────────────────────────────────────────────────────────
+    # 3.  Encoder-only inference util  (FIXED)                       │
+    # ───────────────────────────────────────────────────────────────
+    MASK_ID = tokenizer.mask_token_id or tokenizer.convert_tokens_to_ids("[MASK]")  # <- NEW
+
+    def greedy_pool(idx_order, tag):
         """
-        index_tensor – 1-D tensor of token indices (already on CUDA)
-        which        – "low"  → walk upward
-                       "high" → walk downward
-        Returns (best_pool:list[int], best_acc:float)
+        idx_order : tensor of token-indices sorted hi→lo  or  lo→hi
+        tag       : "high" | "low"   (for the debug print)
+        returns   : (best_pool_indices  ,  best_accuracy)
         """
-        # ---- make everything vanilla Python ints ---------------------
-        indices = index_tensor.tolist()                # e.g. [7, 10, 13, …]
-        if which == "high":
-            indices = indices[::-1]                    # reverse for top-down
-    
-        best_pool: list[int] = []
-        best_acc  = 0.0
-    
-        for i in range(0, len(indices), 2):            # 2 at a time
-            cand   = indices[i : i + 2]                # plain list[int]
-            trial  = best_pool + cand                  # grow pool
-    
-            # ---- build masked input ----------------------------------
-            mask_flags  = torch.ones_like(ids).bool()  # mask everything
-            mask_flags[0, trial] = False               # …except the pool
-            masked_ids  = ids.where(~mask_flags, mask_token_id)
+        best_pool, best_acc = [], 0.0
+        ptr = 0
+        while ptr < len(idx_order):
+            cand = idx_order[ptr : ptr + 2]               # 2-at-a-time
+            pool = best_pool + cand.tolist()              # grow pool
+            ptr += 2
     
-            # ---- second forward-pass ---------------------------------
-            with torch.no_grad():
-                x_m   = emb_drop(emb_ln(embeddings(masked_ids)))
-                ext_m = full_model.bert.get_extended_attention_mask(mask, x_m.shape[:-1])
-                enc_m = encoder(x_m, attention_mask=ext_m)
-                logits = mlm_head(enc_m)[0]            # (S, V)
+            # --- build *mask* for “everything NOT in pool” ----------
+            mask_flags = torch.zeros_like(ids, dtype=torch.bool)
+            mask_flags[0, pool] = True                   # keep these un-masked
+            masked_ids = ids.where(mask_flags, MASK_ID)  # <-  uses the constant
     
-            pred = logits.argmax(-1)
-            corr = (pred[mask_flags] == ids[mask_flags]).float().mean().item()
+            # re-encode & score
+            enc_m   = encode(masked_ids, mask)           # helper already defined
+            logits  = mlm_head(enc_m).logits[0]          # (S, V)
+            preds   = logits.argmax(-1)
+            acc     = (preds[~mask_flags] == ids[0][~mask_flags]).float().mean().item()
     
-            if corr > best_acc:
-                best_acc  = corr
-                best_pool = trial                      # accept improvement
-            if best_acc >= 0.50:
-                break                                  # early exit
+            if acc > best_acc:                           # accept pool only on gain
+                best_pool, best_acc = pool, acc
+                if acc >= 0.50:                          # early-stop rule
+                    break
     
+        print(f"{tag:>4s}-pool  {best_pool}   acc={best_acc:.3f}")
         return best_pool, best_acc