Update app.py

app.py

@@ -35,12 +35,19 @@ def analyze(user_message, assistant_reply):
 
     scores = torch.sigmoid(values[0]).cpu().numpy()
 
-    # Clean up GPT-2 special characters
+    # Only keep tokens that belong to the assistant reply
+    # Find where the assistant reply starts in the token list
+    reply_tokens = tokenizer(assistant_reply, return_tensors="pt").input_ids[0].tolist()
+    n_reply = len(reply_tokens)
+    tokens = tokens[-n_reply:]
+    scores = scores[-n_reply:]
+
     def clean(tok):
         return tok.replace("Ġ", " ").replace("Ċ", "\\n").strip() or tok
 
     labels = [f"{clean(tok)} [{i}]" for i, tok in enumerate(tokens)]
-    df = pd.DataFrame({"token": labels, "value score": scores.tolist()})
+    df = pd.DataFrame({"token": labels, "value score": scores.tolist(), "order": list(range(len(tokens)))})
+    df = df.sort_values("order").drop(columns="order")
 
     stats = (
         f"**Tokens:** {len(tokens)}  |  "