Update app.py

app.py

@@ -90,24 +90,24 @@ def encode_and_trace(text, selected_roles):
     sel_ids = [tokenizer.convert_tokens_to_ids(t) for t in selected_roles]
     sel_ids_tensor = torch.tensor(sel_ids, device="cuda")
 
-    # Tokenize input
+    # Tokenize
     batch = tokenizer(text, return_tensors="pt").to("cuda")
     ids, attn = batch.input_ids, batch.attention_mask
     S = ids.shape[1]
 
-    # Safe encoder forward
+    # Encode helper
     def encode(input_ids, attn_mask):
         x = embeddings(input_ids)
         if emb_ln: x = emb_ln(x)
         if emb_drop: x = emb_drop(x)
         ext = full_model.bert.get_extended_attention_mask(attn_mask, x.shape[:-1])
-        return encoder(x, attention_mask=ext)[0]
+        return encoder(x, attention_mask=ext)[0]  # shape: (1, S, H)
 
     encoded = encode(ids, attn)
 
-    # Get raw symbolic token embeddings directly
-    symbolic_embeds = embeddings.word_embeddings(sel_ids_tensor)  # ✅ FIXED
-    sim = cosine(encoded.unsqueeze(1), symbolic_embeds.unsqueeze(0))  # (S, R)
+    # Project symbolic token embeddings
+    symbolic_embeds = embeddings.word_embeddings(sel_ids_tensor)  # shape: (R, H)
+    sim = cosine(encoded[0].unsqueeze(1), symbolic_embeds.unsqueeze(0))  # (S, R)
     maxcos, argrole = sim.max(-1)  # (S,)
     top_roles = [selected_roles[i] for i in argrole.tolist()]
     sort_idx = maxcos.argsort(descending=True)
@@ -116,7 +116,7 @@ def encode_and_trace(text, selected_roles):
 
     MASK_ID = tokenizer.mask_token_id or tokenizer.convert_tokens_to_ids("[MASK]")
 
-    # 🔧 Pass ids into this function
+    # Final pool evaluator
     def evaluate_pool(idx_order, label, ids):
         best_pool, best_acc = [], 0.0
         ptr = 0
@@ -130,16 +130,17 @@ def encode_and_trace(text, selected_roles):
             masked_input = ids.where(mask_flags, MASK_ID)
 
             encoded_m = encode(masked_input, attn)
-            logits = mlm_head(encoded_m)[0]  # ✅ FIXED — direct tensor
-            preds = logits.argmax(-1)
+            logits = mlm_head(encoded_m)  # (1, S, V)
+            preds = logits.argmax(-1)     # (1, S)
 
-            masked_positions = (~mask_flags[0]).nonzero(as_tuple=False).squeeze(-1)
+            masked_positions = (~mask_flags[0]).nonzero(as_tuple=True)[0]  # 1D tensor
             if masked_positions.numel() == 0:
                 continue
 
-            # ✅ FIXED: indexing from explicitly passed ids
-            gold = ids[0][masked_positions]
-            correct = (preds[masked_positions] == gold).float()
+            # Extract both predicted and gold tokens
+            pred_tokens = preds[0, masked_positions]
+            gold_tokens = ids[0, masked_positions]
+            correct = (pred_tokens == gold_tokens).float()
             acc = correct.mean().item()
 
             if acc > best_acc:
@@ -149,18 +150,18 @@ def encode_and_trace(text, selected_roles):
 
         return best_pool, best_acc
 
-    # Run both pool evaluations
+    # Run both pools
     pool_hi, acc_hi = evaluate_pool(hi_idx, "high", ids)
     pool_lo, acc_lo = evaluate_pool(lo_idx, "low", ids)
 
-    # Per-token symbolic trace
+    # Alignment trace
     decoded_tokens = tokenizer.convert_ids_to_tokens(ids[0])
     role_trace = [
         f"{tok:<15} → {role}   cos={score:.4f}"
         for tok, role, score in zip(decoded_tokens, top_roles, maxcos.tolist())
     ]
 
-    # Output JSON
+    # Return results
     res_json = {
         "High-pool tokens": tokenizer.decode(ids[0, pool_hi]),
         "High accuracy": f"{acc_hi:.3f}",
@@ -174,6 +175,7 @@ def encode_and_trace(text, selected_roles):
 
 
 
+
 # ------------------------------------------------------------------
 # 4.  Gradio UI -----------------------------------------------------
 def build_interface():