Spaces:
Sleeping
Sleeping
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,
|
|
|
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 |
-
|
|
|
|
|
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
|
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")
|