Update app.py

app.py

@@ -113,29 +113,48 @@ def encode_and_trace(text, selected_roles):
     # 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
+    # ------------------------------------------------------------------
+    # Greedy pool helper   (drop-in replacement)
+    # ------------------------------------------------------------------
+    def greedy_pool(index_list, which: str):
+        """
+        index_list  – indices (list[int]) to start pooling from
+        which       – "low"  → walk upward
+                      "high" → walk downward
+        Returns (best_pool, best_acc)
+        """
+        step   = +1 if which == "low" else -1
+        best_pool, best_acc = [], 0.0
+    
+        for i in range(0, len(index_list), 2):          # 2 at a time
+            # current candidate indices to keep un-masked
+            cand = index_list[i : i + 2]
+            pool = best_pool + cand                      # grow pool
+            mask_flags = torch.ones_like(ids).bool()     # mask *everything*
+            mask_flags[0, pool] = False                  # ...except pool
+            masked_ids = ids.masked_fill(~mask_flags, tokenizer.mask_token_id)
+    
+            # ---------- second forward-pass on MASKED input ----------
             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
+                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)          # (1,S,H)
+                logits = mlm_head(enc_m).squeeze(0)                 # (S,V)
+            # ---------------------------------------------------------
+    
+            # accuracy of predicting original tokens only at *masked* positions
+            pred = logits.argmax(-1)
+            corr = (pred[mask_flags] == ids[mask_flags]).float().mean().item()
+    
+            if corr > best_acc:                # greedy improve
+                best_acc   = corr
+                best_pool  = pool
+    
+            # stop early if we already exceed 0.50
+            if best_acc >= 0.50:
+                break
+    
+        return best_pool, best_acc
 
     pool_lo, acc_lo = greedy_pool(lo_idx, "low")
     pool_hi, acc_hi = greedy_pool(hi_idx, "high")