app.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "" #disable cuda

import streamlit as st
import numpy as np
import torch
import time
from Bio import SeqIO

from protxlstm.applications.fitness_prediction import single_mutation_landscape_xlstm, create_mutation_df
from protxlstm.applications.msa_sampler import sample_msa
from protxlstm.models.xlstm import xLSTMLMHeadModel
from protxlstm.utils import load_model
import io

from frontend.constants import info_text, citation_text

if __name__ == "__main__":
    DEFAULT_SEQUENCE = "MTARGLALGLLLLLLCPAQVFSQSCVWYGECGIAYGDKRYNCEYSGPPKPLPKDGYDLVQELCPGFFFGNVSLCCDVRQLQTLKDNLQLPLQFLSRCPSCFYNLLNLFCELTCSPRQSQFLNVTATEDYVDPVTNQTKTNVKELQYYVGQSFANAMYNACRDVEAPSSNDKALGLLCGKDADACNATNWIEYMFNKDNGQAPFTITPVFSDFPVHGMEPMNNATKGCDESVDEVTAPCSCQDCSIVCGPKPQPPPPPAPWTILGLDAMYVIMWITYMAFLLVFFGAFFAVWCYRKRYFVSEYTPIDSNIAFSVNASDKGEASCCDPVSAAFEGCLRRLFTRWGSFCVRNPGCVIFFSLVFITACSSGLVFVRVTTNPVDLWSAPSSQARLEKEYFDQHFGPFFRTEQLIIRAPLTDKHIYQPYPSGADVPFGPPLDIQILHQVLDLQIAIENITASYDNETVTLQDICLAPLSPYNTNCTILSVLNYFQNSHSVLDHKKGDDFFVYADYHTHFLYCVRAPASLNDTSLLHDPCLGTFGGPVFPWLVLGGYDDQNYNNATALVITFPVNNYYNDTEKLQRAQAWEKEFINFVKNYKNPNLTISFTAERSIEDELNRESDSDVFTVVISYAIMFLYISLALGHMKSCRRLLVDSKVSLGIAGILIVLSSVACSLGVFSYIGLPLTLIVIEVIPFLVLAVGVDNIFILVQAYQRDERLQGETLDQQLGRVLGEVAPSMFLSSFSETVAFFLGALSVMPAVHTFSLFAGLAVFIDFLLQITCFVSLLGLDIKRQEKNRLDIFCCVRGAEDGTSVQASESCLFRFFKNSYSPLLLKDWMRPIVIAIFVGVLSFSIAVLNKVDIGLDQSLSMPDDSYMVDYFKSISQYLHAGPPVYFVLEEGHDYTSSKGQNMVCGGMGCNNDSLVQQIFNAAQLDNYTRIGFAPSSWIDDYFDWVKPQSSCCRVDNITDQFCNASVVDPACVRCRPLTPEGKQRPQGGDFMRFLPMFLSDNPNPKCGKGGHAAYSSAVNILLGHGTRVGATYFMTYHTVLQTSADFIDALKKARLIASNVTETMGINGSAYRVFPYSVFYVFYEQYLTIIDDTIFNLGVSLGAIFLVTMVLLGCELWSAVIMCATIAMVLVNMFGVMWLWGISLNAVSLVNLVMSCGISVEFCSHITRAFTVSMKGSRVERAEEALAHMGSSVFSGITLTKFGGIVVLAFAKSQIFQIFYFRMYLAMVLLGATHGLIFLPVLLSYIGPSVNKAKSCATEERYKGTERERLLNF"

    mutation_positions = []
    msa_file = None

    if 'fitness_done' not in st.session_state:
        st.session_state.fitness_done = False
        st.session_state.mutations = None
        st.session_state.fitness_duration = None
        st.session_state.target_sequence = ""
        st.session_state.context_sequences = []
        st.session_state.num_context_sequences = 25

    def run_model():
        try:
            st.session_state.fitness_duration = time.time()
            checkpoint = "protxlstm/checkpoints/small"
            num_context_tokens = 2**15
            df_mutations = create_mutation_df(st.session_state.target_sequence, mutation_positions)
            if msa_file != None and st.session_state.num_context_sequences != 0:
                def load_sequences_from_msa_file(file_obj):
                    text_io = io.TextIOWrapper(file_obj, encoding="utf-8")
                    sequences = [str(record.seq) for record in SeqIO.parse(text_io, "fasta")]
                    return sequences
                msa_sequences = [msa.upper() for msa in load_sequences_from_msa_file(msa_file)]
                st.session_state.context_sequences = sample_msa(msa_sequences, max_context_sequences=st.session_state.num_context_sequences, context_length=num_context_tokens)
                st.session_state.context_sequences += [st.session_state.target_sequence]

            config_update_kwargs = {
                "mlstm_backend": "chunkwise_variable",
                "mlstm_chunksize": 1024,
                "mlstm_return_last_state": True}

            model = load_model(
                checkpoint,
                model_class=xLSTMLMHeadModel,
                device='cpu',
                dtype=torch.bfloat16,
                **config_update_kwargs,
                )
            model = model.eval()
            st.session_state.mutations, _ = single_mutation_landscape_xlstm(model, df_mutations, st.session_state.context_sequences, chunk_chunk_size=2**15)
            print("fitness_done")
            st.session_state.fitness_done = True
            st.session_state.fitness_duration = time.time() - st.session_state.fitness_duration
        except Exception as e:
            print(e)

    # PAGE STYLE (mainly for custom aa selection)
    st.set_page_config(layout="wide")
    st.markdown(
        """
        <style>
        .stButtonGroup button {
            padding: 0px 1px 0px 1px !important;
            border: 0 solid transparent !important;
            min-height: 0px !important;
            line-height: 120% !important;
            height: auto !important;
        }
        .stSidebar {
            width: 600px !important;
        }
        </style>
        """,
        unsafe_allow_html=True
    )


    with st.sidebar:
        st.title("Prot-xLSTM Variant Fitness")

        # LOAD SEQUENCE
        st.session_state.target_sequence = st.text_area(
            "Target protein sequence",
            placeholder=DEFAULT_SEQUENCE,
            value=st.session_state.target_sequence
        )
        if st.button("Load sequence"):
            if st.session_state.target_sequence == "":
                st.session_state.target_sequence = DEFAULT_SEQUENCE

        # MANAGE CONTEXT SEQUENCES
        context_type = st.selectbox(
            "Choose how to enter context",
            ("Enter manually", "Use MSA file"),
            index=None,
            placeholder="Choose context",
        )
        if context_type == 'Enter manually':
            context_sequence_str = st.text_area(
                "Enter context protein sequences (seperated by comma)",
                placeholder=DEFAULT_SEQUENCE,
            )
            st.session_state.context_sequences = context_sequence_str.split(",") + [st.session_state.target_sequence]
            msa_file = None
        elif context_type == 'Use MSA file':
            msa_file = st.file_uploader("Choose MSA file")
            st.session_state.num_context_sequences = st.number_input("How many of these sequences should be used?", min_value=0, step=1, value=25)
        else:
            st.session_state.context_sequences = [st.session_state.target_sequence]
            msa_file = None

    if st.session_state.target_sequence != "":
        with st.container():

            # MUTATION POSITION SELECTION
            aas = list(st.session_state.target_sequence)
            mutation_indices = np.arange(1, len(aas)+1)
            mutation_positions = st.segmented_control(
                "Choose mutation positions (click to select)", mutation_indices, selection_mode="multi", format_func=lambda i: aas[i-1],
            )
            st.button("Check Fitness", on_click=run_model)

            # DISPLAY RESULTS
            if st.session_state.fitness_done:
                st.metric(label="Running time", value=f"{st.session_state.fitness_duration:.2f} sec.")
                selected_pos = st.selectbox(
                    "Visualized mutation position",
                    st.session_state.mutations['position'].unique()
                )
                selected_data = st.session_state.mutations.where(st.session_state.mutations['position'] == selected_pos)
                st.bar_chart(selected_data, x='mutation', y='effect', horizontal=True)
                st.dataframe(st.session_state.mutations, use_container_width=True)

    # TUTORIAL
    with st.expander("Info & Tutorial", expanded=True):
        st.warning('Due to computational constraints, processing may take up to a few minutes.', icon="⚠️")
        st.subheader("Tutorial")
        st.markdown("**1.** Choose a target protein sequence (leave empty to use a sample sequence) and press 'Load Sequence'")
        st.markdown("**2.** Enter or upload you context sequences. (leave empty to use no context)")
        st.markdown("**3.** Choose which amino acids to mutate (click on the AA's to select them) and press 'Check Fitness'")
        st.subheader("General Information")
        st.markdown(info_text, unsafe_allow_html=True)
        st.markdown("")
        st.subheader("Cite us / BibTex")
        st.code(citation_text, language=None)