Upload app.py

app.py

@@ -0,0 +1,193 @@
+import streamlit as st
+import streamlit.components.v1 as components
+import pandas as pd
+import mols2grid
+from ipywidgets import interact, widgets
+import textwrap
+# import numpy as np
+from transformers import EncoderDecoderModel, RobertaTokenizer
+
+from moses.metrics.utils import QED, SA, logP, NP, weight, get_n_rings
+from moses.utils import mapper, get_mol
+
+# @st.cache(allow_output_mutation=False, hash_funcs={Tokenizer: str})
+from typing import List
+
+from util import filter_dataframe
+
+
+@st.cache(suppress_st_warning=True)
+def load_models():
+    # protein_tokenizer = RobertaTokenizer.from_pretrained("gokceuludogan/WarmMolGenTwo")
+    # mol_tokenizer = RobertaTokenizer.from_pretrained("seyonec/PubChem10M_SMILES_BPE_450k")
+    model1 = EncoderDecoderModel.from_pretrained("gokceuludogan/WarmMolGenOne")
+    model2 = EncoderDecoderModel.from_pretrained("gokceuludogan/WarmMolGenTwo")
+    return model1, model2  # , protein_tokenizer, mol_tokenizer
+
+
+def count(smiles_list: List[str]):
+    counts = []
+    for smiles in smiles_list:
+        counts.append(len(smiles))
+
+    return counts
+
+
+def remove_none_elements(mol_list, smiles_list):
+    filtered_mol_list = []
+    filtered_smiles_list = []
+    indices = []
+    for i, element in enumerate(mol_list):
+        if element is not None:
+            filtered_mol_list.append(element)
+        else:
+            indices.append(i)
+    removed_len = len(indices)
+
+    for i in range(len(smiles_list)):
+        if i not in indices:
+            filtered_smiles_list.append(smiles_list.__getitem__(i))
+
+    return filtered_mol_list, filtered_smiles_list, removed_len
+
+
+def format_list_numbers(lst):
+    for i, value in enumerate(lst):
+        lst[i] = float("{:.3f}".format(value))
+
+
+def generate_molecules(model_name, num_mols, max_new_tokens, do_sample, num_beams, target, pool):
+    protein_tokenizer = RobertaTokenizer.from_pretrained("gokceuludogan/WarmMolGenTwo")
+    mol_tokenizer = RobertaTokenizer.from_pretrained("seyonec/PubChem10M_SMILES_BPE_450k")
+    # model1, model2, protein_tokenizer, mol_tokenizer = load_models()
+    model1, model2 = load_models()
+    inputs = protein_tokenizer(target, return_tensors="pt")
+
+    model = model1 if model_name == 'WarmMolGenOne' else model2
+    outputs = model.generate(**inputs, decoder_start_token_id=mol_tokenizer.bos_token_id,
+                             eos_token_id=mol_tokenizer.eos_token_id, pad_token_id=mol_tokenizer.eos_token_id,
+                             max_length=int(max_new_tokens), num_return_sequences=int(num_mols),
+                             do_sample=do_sample, num_beams=num_beams)
+    output_smiles = mol_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    st.write("### Generated Molecules")
+    # mol_list = list(map(MolFromSmiles, output_smiles))
+    # print(mol_list)
+    # QED_scores = list(map(QED.qed, mol_list))
+    # print(QED_scores)
+    # st.write(output_smiles)
+    mol_list = mapper(pool)(get_mol, output_smiles)
+    mol_list, output_smiles, removed_len = remove_none_elements(mol_list, output_smiles)
+    if removed_len != 0:
+        st.write(f"#### Note that: {removed_len} numbers of generated invalid molecules are discarded.")
+
+    QED_scores = mapper(pool)(QED, mol_list)
+    SA_scores = mapper(pool)(SA, mol_list)
+    logP_scores = mapper(pool)(logP, mol_list)
+    NP_scores = mapper(pool)(NP, mol_list)
+    weight_scores = mapper(pool)(weight, mol_list)
+
+    format_list_numbers(QED_scores)
+    format_list_numbers(SA_scores)
+    format_list_numbers(logP_scores)
+    format_list_numbers(NP_scores)
+    format_list_numbers(weight_scores)
+
+    df_smiles = pd.DataFrame(
+        {'SMILES': output_smiles, "QED": QED_scores, "SA": SA_scores, "logP": logP_scores, "NP": NP_scores,
+         "Weight": weight_scores})
+
+    return df_smiles
+
+
+def warm_molgen_demo():
+    with st.form("my_form"):
+        with st.sidebar:
+            st.sidebar.subheader("Configurable parameters")
+
+            model_name = st.sidebar.selectbox(
+                "Model Selector",
+                options=[
+                    "WarmMolGenOne",
+                    "WarmMolGenTwo",
+                ],
+                index=0,
+            )
+
+            num_mols = st.sidebar.number_input(
+                "Number of generated molecules",
+                min_value=0,
+                max_value=20,
+                value=10,
+                help="The number of molecules to be generated.",
+            )
+
+            max_new_tokens = st.sidebar.number_input(
+                "Maximum length",
+                min_value=0,
+                max_value=1024,
+                value=128,
+                help="The maximum length of the sequence to be generated.",
+            )
+            do_sample = st.sidebar.selectbox(
+                "Sampling?",
+                (True, False),
+                help="Whether or not to use sampling; use beam decoding otherwise.",
+            )
+            target = st.text_area(
+                "Target Sequence",
+                "MENTENSVDSKSIKNLEPKIIHGSESMDSGISLDNSYKMDYPEMGLCIIINNKNFHKSTG",
+            )
+            generate_new_molecules = st.form_submit_button("Generate Molecules")
+
+    num_beams = None if do_sample is True else int(num_mols)
+
+    pool = 1
+
+    if generate_new_molecules:
+        st.session_state.df = generate_molecules(model_name, num_mols, max_new_tokens, do_sample, num_beams,
+                                                 target, pool)
+    if 'df' not in st.session_state:
+        st.session_state.df = generate_molecules(model_name, num_mols, max_new_tokens, do_sample, num_beams,
+                                                 target, pool)
+    df = st.session_state.df
+
+    filtered_df = filter_dataframe(df)
+    if filtered_df.empty:
+        st.markdown(
+            """
+            <span style='color: blue; font-size: 30px;'>No molecules were found with specified properties.</span>
+        """,
+            unsafe_allow_html=True
+        )
+    else:
+        raw_html = mols2grid.display(filtered_df, height=1000)._repr_html_()
+        components.html(raw_html, width=900, height=450, scrolling=True)
+
+    st.markdown("## How to Generate")
+    generation_code = f"""
+    from transformers import EncoderDecoderModel, RobertaTokenizer
+    protein_tokenizer = RobertaTokenizer.from_pretrained("gokceuludogan/{model_name}")
+    mol_tokenizer = RobertaTokenizer.from_pretrained("seyonec/PubChem10M_SMILES_BPE_450k")
+    model = EncoderDecoderModel.from_pretrained("gokceuludogan/{model_name}")
+    inputs = protein_tokenizer("{target}", return_tensors="pt")
+    outputs = model.generate(**inputs, decoder_start_token_id=mol_tokenizer.bos_token_id,
+                             eos_token_id=mol_tokenizer.eos_token_id, pad_token_id=mol_tokenizer.eos_token_id,
+                             max_length={max_new_tokens}, num_return_sequences={num_mols}, do_sample={do_sample}, num_beams={num_beams})
+    mol_tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    """
+    st.code(textwrap.dedent(generation_code))  # textwrap.dedent("".join("Halletcez")))
+
+
+st.set_page_config(page_title="WarmMolGen Demo", page_icon="🔥", layout='wide')
+st.markdown("# WarmMolGen Demo")
+st.sidebar.header("WarmMolGen Demo")
+st.markdown(
+    """
+    This demo illustrates WarmMolGen models' generation capabilities.
+    Given a target sequence and a set of parameters, the models generate molecules targeting the given protein sequence.
+    Please enter an input sequence below 👇  and configure parameters from the sidebar 👈 to generate molecules!
+    See below for saving the output molecules and the code snippet generating them!
+"""
+)
+
+warm_molgen_demo()