aksell commited on
Commit
2d68959
·
1 Parent(s): 48a0092

Store selected head and layer sequence

Browse files
hexviz/pages/🗺️Identify_Interesting_Heads.py CHANGED
@@ -3,7 +3,8 @@ import streamlit as st
3
  from hexviz.attention import get_attention, get_sequence, get_structure
4
  from hexviz.models import Model, ModelType
5
  from hexviz.plot import plot_tiled_heatmap
6
- from hexviz.view import menu_items, select_model, select_pdb
 
7
 
8
  st.set_page_config(layout="wide", menu_items=menu_items)
9
  st.subheader("Find interesting heads and layers")
@@ -34,17 +35,15 @@ sequence = get_sequence(selected_chain)
34
 
35
  l = len(sequence)
36
  st.sidebar.markdown("Sequence segment to plot")
37
- slice_start, slice_end = st.sidebar.slider("Sequence", min_value=1, max_value=l, value=(1, 50), step=1)
 
 
38
  # slice_start= st.sidebar.number_input(f"Section start(1-{l})",value=1, min_value=1, max_value=l)
39
  # slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1, max_value=l)
40
  truncated_sequence = sequence[slice_start-1:slice_end]
41
 
42
- head_range = st.sidebar.slider("Heads to plot", min_value=1, max_value=selected_model.heads, value=(1, selected_model.heads//2), step=1)
43
- layer_range = st.sidebar.slider("Layers to plot", min_value=1, max_value=selected_model.layers, value=(1, selected_model.layers//2), step=1)
44
- step_size = st.sidebar.number_input("Optional step size to skip heads and layers", value=1, min_value=1, max_value=selected_model.layers)
45
- layer_sequence = list(range(layer_range[0]-1, layer_range[1], step_size))
46
- head_sequence = list(range(head_range[0]-1, head_range[1], step_size))
47
 
 
48
 
49
  st.markdown(f"Each tile is a heatmap of attention for a section of {pdb_id}({chain_selection}) from residue {slice_start} to {slice_end}. Adjust the section length and starting point in the sidebar.")
50
 
 
3
  from hexviz.attention import get_attention, get_sequence, get_structure
4
  from hexviz.models import Model, ModelType
5
  from hexviz.plot import plot_tiled_heatmap
6
+ from hexviz.view import (menu_items, select_heads_and_layers, select_model,
7
+ select_pdb)
8
 
9
  st.set_page_config(layout="wide", menu_items=menu_items)
10
  st.subheader("Find interesting heads and layers")
 
35
 
36
  l = len(sequence)
37
  st.sidebar.markdown("Sequence segment to plot")
38
+ if "sequence-slice" not in st.session_state:
39
+ st.session_state["sequence-slice"] = (1, min(50, l))
40
+ slice_start, slice_end = st.sidebar.slider("Sequence", key="sequence-slice", min_value=1, max_value=l, step=1)
41
  # slice_start= st.sidebar.number_input(f"Section start(1-{l})",value=1, min_value=1, max_value=l)
42
  # slice_end = st.sidebar.number_input(f"Section end(1-{l})",value=50, min_value=1, max_value=l)
43
  truncated_sequence = sequence[slice_start-1:slice_end]
44
 
 
 
 
 
 
45
 
46
+ layer_sequence, head_sequence = select_heads_and_layers(st.sidebar, selected_model)
47
 
48
  st.markdown(f"Each tile is a heatmap of attention for a section of {pdb_id}({chain_selection}) from residue {slice_start} to {slice_end}. Adjust the section length and starting point in the sidebar.")
49
 
hexviz/view.py CHANGED
@@ -19,8 +19,15 @@ def select_model(models):
19
 
20
  Saves the selected model in session storage.
21
  """
 
22
  selected_model_name = st.selectbox("Select model", [model.name.value for model in models], index=get_selecte_model_index(models))
23
  st.session_state.selected_model_name = selected_model_name
 
 
 
 
 
 
24
  select_model = next((model for model in models if model.name.value == selected_model_name), None)
25
  return select_model
26
 
@@ -38,4 +45,28 @@ def select_pdb():
38
  st.session_state.selected_chains = None
39
  st.session_state.selected_chain_index = 0
40
  st.session_state.pdb_id = pdb_id
41
- return pdb_id
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
20
  Saves the selected model in session storage.
21
  """
22
+ stored_model = st.session_state.get("selected_model_name", None)
23
  selected_model_name = st.selectbox("Select model", [model.name.value for model in models], index=get_selecte_model_index(models))
24
  st.session_state.selected_model_name = selected_model_name
25
+ model_changed = stored_model != selected_model_name
26
+ if model_changed:
27
+ if "plot_heads" in st.session_state:
28
+ del st.session_state.plot_heads
29
+ if "plot_layers" in st.session_state:
30
+ del st.session_state.plot_layers
31
  select_model = next((model for model in models if model.name.value == selected_model_name), None)
32
  return select_model
33
 
 
45
  st.session_state.selected_chains = None
46
  st.session_state.selected_chain_index = 0
47
  st.session_state.pdb_id = pdb_id
48
+ return pdb_id
49
+
50
+ def select_heads_and_layers(sidebar, model):
51
+ sidebar.markdown(
52
+ """
53
+ Select Heads and Layers
54
+ ---
55
+ """
56
+ )
57
+ if "plot_heads" not in st.session_state:
58
+ "setting heads"
59
+ st.session_state.plot_heads = (1, model.heads//2)
60
+ if "plot_layers" not in st.session_state:
61
+ "setting layers"
62
+ st.session_state.plot_layers = (1, model.layers//2)
63
+ head_range = sidebar.slider("Heads to plot", min_value=1, max_value=model.heads, value=st.session_state.plot_heads, step=1)
64
+ st.session_state.plot_heads = head_range
65
+ layer_range = sidebar.slider("Layers to plot", min_value=1, max_value=model.layers, value=st.session_state.plot_layers, step=1)
66
+ st.session_state.plot_layers = layer_range
67
+
68
+ step_size = sidebar.number_input("Optional step size to skip heads and layers", value=1, min_value=1, max_value=model.layers)
69
+ layer_sequence = list(range(layer_range[0]-1, layer_range[1], step_size))
70
+ head_sequence = list(range(head_range[0]-1, head_range[1], step_size))
71
+
72
+ return layer_sequence, head_sequence
hexviz/🧬Attention_Visualization.py CHANGED
@@ -6,7 +6,7 @@ from stmol import showmol
6
 
7
  from hexviz.attention import get_attention_pairs, get_chains, get_structure
8
  from hexviz.models import Model, ModelType
9
- from hexviz.view import select_model, select_pdb, menu_items
10
 
11
  st.set_page_config(layout="centered", menu_items=menu_items)
12
  st.title("Attention Visualization on proteins")
 
6
 
7
  from hexviz.attention import get_attention_pairs, get_chains, get_structure
8
  from hexviz.models import Model, ModelType
9
+ from hexviz.view import menu_items, select_model, select_pdb
10
 
11
  st.set_page_config(layout="centered", menu_items=menu_items)
12
  st.title("Attention Visualization on proteins")