Update app.py

app.py

@@ -90,7 +90,6 @@ def encode_and_trace(text: str, selected_roles: list[str]):
         ids, mask = batch.input_ids, batch.attention_mask
 
         x = emb_drop(emb_ln(embeddings(ids)))
-
         ext_mask = full_model.bert.get_extended_attention_mask(mask, x.shape[:-1])
         enc = encoder(x, attention_mask=ext_mask)              # (1, S, H)
 
@@ -98,20 +97,20 @@ def encode_and_trace(text: str, selected_roles: list[str]):
         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]
+        found_tokens = [tokenizer.convert_ids_to_tokens([tid])[0]
+                        for tid in ids[0].tolist() if tid in sel_ids]
+        tokens_str = ", ".join(found_tokens) or "(none)"
 
         if flags.any():
-            vec = enc[0][flags].mean(0)
+            vec  = enc[0][flags].mean(0)
             norm = f"{vec.norm().item():.4f}"
         else:
             norm = "0.0000"
 
-        return {
-            "Symbolic Tokens": ", ".join(found) or "(none)",
-            "Embedding Norm":  norm,
-            "Symbolic Token Count": int(flags.sum().item()),
-        }
+        count = int(flags.sum().item())
+        # >>> return *three* scalars, not one dict  <<<
+        return tokens_str, norm, count
+
 
 
 # ------------------------------------------------------------------