main script statement added

app.py

@@ -15,137 +15,138 @@ import io
 
 from frontend.constants import info_text, citation_text
 
-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 __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
     )
-    if context_type == 'Enter manually':
-        context_sequence_str = st.text_area(
-            "Enter context protein sequences (seperated by comma)",
+
+
+    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
         )
-        st.session_state.context_sequences = context_sequence_str.split(",") + [st.session_state.target_sequence]
-    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]
-
-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],
+        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",
         )
-        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()
+        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]
+        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]
+
+    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],
             )
-            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.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)
+            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.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)