Fix x-axis bug

app.py

@@ -114,7 +114,8 @@ def imshow(tensor, xlabel="X", ylabel="Y", zlabel=None, xticks=None, yticks=None
 
 def plot_residual_stream_patch(clean_prompt=None, answer=None, corrupt_prompt=None, corrupt_answer=None):
     layers = ["blocks.0.hook_resid_pre", *[f"blocks.{i}.hook_resid_post" for i in range(model.cfg.n_layers)]]
-    token_labels = model.to_str_tokens(clean_prompt)
+    clean_tokens = model.to_str_tokens(clean_prompt)
+    token_labels = [f"(pos {i:2}) {t}" for i, t in enumerate(clean_tokens)]
     patching_effect = compute_residual_stream_patch(clean_prompt=clean_prompt, answer=answer, corrupt_prompt=corrupt_prompt, corrupt_answer=corrupt_answer, layers=layers)
     fig = imshow(patching_effect, xticks=token_labels, yticks=layers, xlabel="Position", ylabel="Layer",
        zlabel="Logit Difference", title="Patching residual stream at specific layer and position")