Update app.py

app.py

@@ -121,7 +121,7 @@ def encode_and_trace(text: str, selected_roles: list[str]):
         rows = []
         for role in selected_roles:
             rid   = tokenizer.convert_tokens_to_ids(role)
-            rvec  = embeddings.weight[rid]                       # (H,)
+            rvec  = embeddings.word_embeddings.weight[rid]                    # (H,)
             # cosine similarity to every *input* token embedding
             sims  = F.cosine_similarity(rvec.unsqueeze(0), token_emb, dim=-1)
             best  = torch.argmax(sims).item()