Update app.py

app.py

@@ -112,50 +112,53 @@ def encode_and_trace(text, selected_roles):
 
     # container for summary text
     report_lines = []
-
+    
     # ------------------------------------------------------------------
-    # Greedy pool helper   (drop-in replacement)
+    # Greedy pool helper  – tensor-safe version
     # ------------------------------------------------------------------
-    def greedy_pool(index_list, which: str):
+    def greedy_pool(index_tensor: torch.Tensor, which: str):
         """
-        index_list  – indices (list[int]) to start pooling from
-        which       – "low"  → walk upward
-                      "high" → walk downward
-        Returns (best_pool, best_acc)
+        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)
         """
-        step   = +1 if which == "low" else -1
-        best_pool, best_acc = [], 0.0
+        # ---- 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(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)
+        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
     
-            # ---------- second forward-pass on MASKED input ----------
+            # ---- 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)
+    
+            # ---- 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)          # (1,S,H)
-                logits = mlm_head(enc_m).squeeze(0)                 # (S,V)
-            # ---------------------------------------------------------
+                enc_m = encoder(x_m, attention_mask=ext_m)
+                logits = mlm_head(enc_m)[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 corr > best_acc:
+                best_acc  = corr
+                best_pool = trial                      # accept improvement
             if best_acc >= 0.50:
-                break
+                break                                  # early exit
     
         return best_pool, best_acc
 
+
     pool_lo, acc_lo = greedy_pool(lo_idx, "low")
     pool_hi, acc_hi = greedy_pool(hi_idx, "high")