update illustrations of respective encoders

app.py

@@ -56,85 +56,98 @@ def predict_dti():
     
     with col1:
         st.markdown('### Drug')
-        smiles = st.text_input('Enter the SMILES of the query drug compound', value='CC(=O)OC1=CC=CC=C1C(=O)O', placeholder='CC(=O)OC1=CC=CC=C1C(=O)O')
         
-        if smiles:
-            mol = Chem.MolFromSmiles(smiles)
-            mol_img = Chem.Draw.MolToImage(mol)
-            st.image(mol_img) #, width = 140)
-            
+        mol_col1, mol_col2 = st.columns(2)
+
+        with mol_col1:
+            smiles = st.text_input('Enter the SMILES of the query drug compound', value='CC(=O)OC1=CC=CC=C1C(=O)O', placeholder='CC(=O)OC1=CC=CC=C1C(=O)O')
+            if smiles:
+                mol = Chem.MolFromSmiles(smiles)
+                mol_img = Chem.Draw.MolToImage(mol)
+                st.image(mol_img) #, width = 140)
+
+        with mol_col2:
             selected_encoder = st.selectbox(
                 'Select encoder for drug compound',('None', 'CDDD', 'MolBERT')
             )
-            if selected_encoder == 'CDDD':
-                from cddd.inference import InferenceModel
-                CDDD_MODEL_DIR = 'src/encoders/cddd'
-                cddd_model = InferenceModel(CDDD_MODEL_DIR)
-                embedding = cddd_model.seq_to_emb([smiles])
-                #from huggingface_hub import hf_hub_download
-                #precomputed_embs = f'{selected_encoder}_encoding.csv'
-                #REPO_ID = "emmas96/Lenselink"
-                #embs_path = hf_hub_download(REPO_ID, precomputed_embs)
-                #embs = pd.read_csv(embs_path)
-                #embedding = embs[smiles]
-            elif selected_encoder == 'MolBERT':
-                from molbert.utils.featurizer.molbert_featurizer import MolBertFeaturizer
-                from huggingface_hub import hf_hub_download
-                CDDD_MODEL_DIR = 'encoders/molbert/last.ckpt'
-                REPO_ID = "emmas96/hyperpcm"
-                checkpoint_path = hf_hub_download(REPO_ID, MOLBERT_MODEL_DIR)
-                molbert_model = MolBertFeaturizer(checkpoint_path, max_seq_len=500, embedding_type='average-1-cat-pooled')
-                embedding = molbert_model.transform([smiles])
-            else: 
-                st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
-                embedding = None
-            if embedding is not None:
-                st.write(f'{selected_encoder} embedding')
-                st.write(embedding)
+            st.image('molecule_encoder.png')
+            if smiles: 
+                if selected_encoder == 'CDDD':
+                    from cddd.inference import InferenceModel
+                    CDDD_MODEL_DIR = 'src/encoders/cddd'
+                    cddd_model = InferenceModel(CDDD_MODEL_DIR)
+                    embedding = cddd_model.seq_to_emb([smiles])
+                    #from huggingface_hub import hf_hub_download
+                    #precomputed_embs = f'{selected_encoder}_encoding.csv'
+                    #REPO_ID = "emmas96/Lenselink"
+                    #embs_path = hf_hub_download(REPO_ID, precomputed_embs)
+                    #embs = pd.read_csv(embs_path)
+                    #embedding = embs[smiles]
+                elif selected_encoder == 'MolBERT':
+                    from molbert.utils.featurizer.molbert_featurizer import MolBertFeaturizer
+                    from huggingface_hub import hf_hub_download
+                    CDDD_MODEL_DIR = 'encoders/molbert/last.ckpt'
+                    REPO_ID = "emmas96/hyperpcm"
+                    checkpoint_path = hf_hub_download(REPO_ID, MOLBERT_MODEL_DIR)
+                    molbert_model = MolBertFeaturizer(checkpoint_path, max_seq_len=500, embedding_type='average-1-cat-pooled')
+                    embedding = molbert_model.transform([smiles])
+                else: 
+                    st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
+                    embedding = None
+                if embedding is not None:
+                    st.write(f'{selected_encoder} embedding')
+                    st.write(embedding)
 
     with col2:
         st.markdown('### Target')
-        sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
+
+        prot_col1, prot_col2 = st.columns(2)
+
+        with prot_col1:
+            sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
         
-        if sequence:
-            st.markdown('\n\n\n\n Plot of protein to be added soon. \n\n\n\n')
-            
+            if sequence:
+                st.markdown('\n\n\n\n Plot of protein to be added soon. \n\n\n\n')
+
+        with prot_col2:
             selected_encoder = st.selectbox(
                 'Select encoder for protein target',('None', 'SeqVec', 'UniRep', 'ESM-1b', 'ProtT5')
             )
-            if selected_encoder == 'SeqVec':
-                from bio_embeddings.embed import SeqVecEmbedder
-                encoder = SeqVecEmbedder()
-                embeddings = encoder.embed_batch([sequence])
-                for emb in embeddings:
-                    embedding = encoder.reduce_per_protein(emb)
-                    break
-            elif selected_encoder == 'UniRep':
-                from jax_unirep.utils import load_params
-                params = load_params() 
-                from jax_unirep.featurize import get_reps                               
-                embedding, h_final, c_final = get_reps([sequence])
-                embedding = embedding.mean(axis=0)
-            elif selected_encoder == 'ESM-1b':
-                from bio_embeddings.embed import ESM1bEmbedder
-                encoder = ESM1bEmbedder()
-                embeddings = encoder.embed_batch([sequence])
-                for emb in embeddings:
-                    embedding = encoder.reduce_per_protein(emb)
-                    break
-            elif selected_encoder == 'ProtT5':
-                from bio_embeddings.embed import ProtTransT5XLU50Embedder
-                encoder = ProtTransT5XLU50Embedder()
-                embeddings = encoder.embed_batch([sequence])
-                for emb in embeddings:
-                    embedding = encoder.reduce_per_protein(emb)
-                    break
-            else: 
-                st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
-                embedding = None
-            if embedding is not None:
-                st.write(f'{selected_encoder} embedding')
-                st.write(embedding)
+            st.image('protein_encoder.png')
+            if sequence: 
+                if selected_encoder == 'SeqVec':
+                    from bio_embeddings.embed import SeqVecEmbedder
+                    encoder = SeqVecEmbedder()
+                    embeddings = encoder.embed_batch([sequence])
+                    for emb in embeddings:
+                        embedding = encoder.reduce_per_protein(emb)
+                        break
+                elif selected_encoder == 'UniRep':
+                    from jax_unirep.utils import load_params
+                    params = load_params() 
+                    from jax_unirep.featurize import get_reps                               
+                    embedding, h_final, c_final = get_reps([sequence])
+                    embedding = embedding.mean(axis=0)
+                elif selected_encoder == 'ESM-1b':
+                    from bio_embeddings.embed import ESM1bEmbedder
+                    encoder = ESM1bEmbedder()
+                    embeddings = encoder.embed_batch([sequence])
+                    for emb in embeddings:
+                        embedding = encoder.reduce_per_protein(emb)
+                        break
+                elif selected_encoder == 'ProtT5':
+                    from bio_embeddings.embed import ProtTransT5XLU50Embedder
+                    encoder = ProtTransT5XLU50Embedder()
+                    embeddings = encoder.embed_batch([sequence])
+                    for emb in embeddings:
+                        embedding = encoder.reduce_per_protein(emb)
+                        break
+                else: 
+                    st.write('No pre-trained version of HyperPCM is available for the chosen encoder.')
+                    embedding = None
+                if embedding is not None:
+                    st.write(f'{selected_encoder} embedding')
+                    st.write(embedding)
 
     st.write('TODO run inference with HyperPCM on the given drug compound and protein target.')
     
@@ -148,7 +161,7 @@ def retrieval():
     sequence = st.text_input('Enter the amino-acid sequence of the query protein target', value='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA', placeholder='HXHVWPVQDAKARFSEFLDACITEGPQIVSRRGAEEAVLVPIGEWRRLQAAA')
     
     if sequence:
-        col1, col2, col3 = st.columns(3)
+        col1, col2 = st.columns(2)
         with col1:
             st.markdown('\n\n\n\n Plot of protein to be added soon. \n\n\n\n')
         
@@ -162,8 +175,7 @@ def retrieval():
             for emb in embeddings:
                 embedding = encoder.reduce_per_protein(emb)
                 break
-        with col3:
-            st.write(f'SeqVec embedding')
+            st.write('SeqVec embedding')
             st.write(embedding)
             st.write(np.transpose(embedding))