Spaces:
Sleeping
Sleeping
| from io import StringIO | |
| import streamlit as st | |
| from Bio.PDB import PDBParser | |
| from hexviz.attention import get_pdb_file, get_pdb_from_seq | |
| menu_items = { | |
| "Get Help": "https://huggingface.co/spaces/aksell/hexviz/discussions/new", | |
| "Report a bug": "https://huggingface.co/spaces/aksell/hexviz/discussions/new", | |
| "About": "Created by [Aksel Lenes](https://github.com/aksell/) from Noelia Ferruz's group at the Institute of Molecular Biology of Barcelona. Read more at https://www.aiproteindesign.com/", | |
| } | |
| def get_selecte_model_index(models): | |
| selected_model_name = st.session_state.get("selected_model_name", None) | |
| if selected_model_name is None: | |
| return 0 | |
| else: | |
| return next( | |
| (i for i, model in enumerate(models) if model.name.value == selected_model_name), | |
| None, | |
| ) | |
| def clear_model_state(): | |
| if "plot_heads" in st.session_state: | |
| del st.session_state.plot_heads | |
| if "plot_layers" in st.session_state: | |
| del st.session_state.plot_layers | |
| if "selected_head" in st.session_state: | |
| del st.session_state.selected_head | |
| if "selected_layer" in st.session_state: | |
| del st.session_state.selected_layer | |
| if "plot_layers" in st.session_state: | |
| del st.session_state.plot_layers | |
| if "plot_heads" in st.session_state: | |
| del st.session_state.plot_heads | |
| if "label_tokens" in st.session_state: | |
| del st.session_state.label_tokens | |
| def select_model(models): | |
| if "selected_model_name" not in st.session_state: | |
| st.session_state.selected_model_name = models[0].name.value | |
| selected_model_name = st.selectbox( | |
| "Select model", | |
| [model.name.value for model in models], | |
| key="selected_model_name", | |
| on_change=clear_model_state, | |
| ) | |
| select_model = next( | |
| (model for model in models if model.name.value == selected_model_name), None | |
| ) | |
| return select_model | |
| def clear_pdb_state(): | |
| if "selected_chains" in st.session_state: | |
| del st.session_state.selected_chains | |
| if "selected_chain" in st.session_state: | |
| del st.session_state.selected_chain | |
| if "sequence_slice" in st.session_state: | |
| del st.session_state.sequence_slice | |
| if "uploaded_pdb_str" in st.session_state: | |
| del st.session_state.uploaded_pdb_str | |
| def select_pdb(): | |
| if "pdb_id" not in st.session_state: | |
| st.session_state.pdb_id = "2FZ5" | |
| pdb_id = st.text_input(label="1.PDB ID", key="pdb_id", on_change=clear_pdb_state) | |
| return pdb_id | |
| def select_protein(pdb_code, uploaded_file, input_sequence): | |
| # We get the pdb from 1 of 3 places: | |
| # 1. Cached pdb from session storage | |
| # 2. PDB file from uploaded file | |
| # 3. PDB file fetched based on the pdb_code input | |
| parser = PDBParser() | |
| if uploaded_file is not None: | |
| pdb_str = uploaded_file.read().decode("utf-8") | |
| st.session_state["uploaded_pdb_str"] = pdb_str | |
| source = f"uploaded pdb file {uploaded_file.name}" | |
| structure = parser.get_structure("Userfile", StringIO(pdb_str)) | |
| elif input_sequence: | |
| pdb_str = get_pdb_from_seq(str(input_sequence)) | |
| if not pdb_str: | |
| st.error("ESMfold error, unable to fold sequence") | |
| return None, None, None | |
| else: | |
| structure = parser.get_structure("ESMFold", StringIO(pdb_str)) | |
| if "selected_chains" in st.session_state: | |
| del st.session_state.selected_chains | |
| source = "Input sequence + ESM-fold" | |
| elif "uploaded_pdb_str" in st.session_state: | |
| pdb_str = st.session_state.uploaded_pdb_str | |
| source = "Uploaded file stored in cache" | |
| structure = parser.get_structure("userfile", StringIO(pdb_str)) | |
| else: | |
| file = get_pdb_file(pdb_code) | |
| pdb_str = file.read() | |
| source = f"PDB ID: {pdb_code}" | |
| structure = parser.get_structure(pdb_code, StringIO(pdb_str)) | |
| return pdb_str, structure, source | |
| def select_heads_and_layers(sidebar, model): | |
| sidebar.markdown( | |
| """ | |
| Select Heads and Layers | |
| --- | |
| """ | |
| ) | |
| if "plot_heads" not in st.session_state: | |
| st.session_state.plot_heads = (1, model.heads // 2) | |
| head_range = sidebar.slider( | |
| "Heads to plot", min_value=1, max_value=model.heads, key="plot_heads", step=1 | |
| ) | |
| if "plot_layers" not in st.session_state: | |
| st.session_state.plot_layers = (1, model.layers // 2) | |
| layer_range = sidebar.slider( | |
| "Layers to plot", min_value=1, max_value=model.layers, key="plot_layers", step=1 | |
| ) | |
| if "plot_step_size" not in st.session_state: | |
| st.session_state.plot_step_size = 1 | |
| step_size = sidebar.number_input( | |
| "Optional step size to skip heads and layers", | |
| key="plot_step_size", | |
| min_value=1, | |
| max_value=model.layers, | |
| ) | |
| layer_sequence = list(range(layer_range[0] - 1, layer_range[1], step_size)) | |
| head_sequence = list(range(head_range[0] - 1, head_range[1], step_size)) | |
| return layer_sequence, head_sequence | |
| def select_sequence_slice(sequence_length): | |
| st.sidebar.markdown( | |
| """ | |
| Sequence segment to plot | |
| --- | |
| """ | |
| ) | |
| if "sequence_slice" not in st.session_state: | |
| st.session_state.sequence_slice = (1, min(50, sequence_length)) | |
| slice = st.sidebar.slider( | |
| "Sequence", key="sequence_slice", min_value=1, max_value=sequence_length, step=1 | |
| ) | |
| return slice | |