Bug: Label hidden tokens for ProtT5 too

hexviz/pages/1_🗺️Identify_Interesting_Heads.py

@@ -141,7 +141,8 @@ with right:
 
 if selected_model.name == ModelType.PROT_T5:
     # Remove leading underscores from residue tokens
-    tokens = [token[1:] for token in tokens if str(token) != "</s>"]
+    tokens = [token[1:] if str(token) != "</s>" else token for token in tokens]
+
 
 single_head_fig = plot_single_heatmap(attention, layer, head, tokens=tokens)
 st.pyplot(single_head_fig)