Fix bug in column vise attention calc

hexviz/attention.py

@@ -293,9 +293,9 @@ def get_attention_pairs(
                             residue_attention.get(res - ec_tag_length, 0) + attn_value
                         )
         if not ec_number:
-            attention_into_res = attention[head, layer].sum(dim=0)
+            attention_into_res = attention[layer, head].sum(dim=0)
         else:
-            attention_into_res = attention[head, layer, ec_tag_length:, ec_tag_length:].sum(dim=0)
+            attention_into_res = attention[layer, head, ec_tag_length:, ec_tag_length:].sum(dim=0)
         top_n_values, top_n_indexes = torch.topk(attention_into_res, top_n)
 
         for res, attn_sum in zip(top_n_indexes, top_n_values):

hexviz/🧬Attention_Visualization.py

@@ -28,7 +28,7 @@ models = [
 ]
 
 with st.expander("Input a PDB id, upload a PDB file or input a sequence", expanded=True):
-    pdb_id = select_pdb() or "2WK4"
+    pdb_id = select_pdb()
     uploaded_file = st.file_uploader("2.Upload PDB", type=["pdb"])
     input_sequence = st.text_area("3.Input sequence", "", key="input_sequence", max_chars=400)
     sequence, error = clean_and_validate_sequence(input_sequence)