add dummy retrieval page

app.py

@@ -136,6 +136,85 @@ def display_dti():
                 st.write(f'{selected_encoder} embedding')
                 st.write(embedding)
 
+def retrieval():
+    st.markdown('##')
+
+    st.markdown('### Target')
+    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')
+        
+        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.markdown('### Retrieval of top-5 drug coupound from ChEMBL:')
+    col1, col2, col3, col4, col5 = st.columns(5)
+    with col1:
+        smiles = 'CC(=O)OC1=CC=CC=C1C(=O)O'
+        mol = Chem.MolFromSmiles(smiles)
+        mol_img = Chem.Draw.MolToImage(mol)
+        st.image(mol_img) 
+
+    with col2:
+        smiles = 'COc1cc(C=O)ccc1O'
+        mol = Chem.MolFromSmiles(smiles)
+        mol_img = Chem.Draw.MolToImage(mol)
+        st.image(mol_img) 
+
+    with col3:
+        smiles = 'CC(=O)Nc1ccc(O)cc1'
+        mol = Chem.MolFromSmiles(smiles)
+        mol_img = Chem.Draw.MolToImage(mol)
+        st.image(mol_img) 
+
+    with col4:
+        smiles = 'CC(=O)Nc1ccc(OS(=O)(=O)O)cc1'
+        mol = Chem.MolFromSmiles(smiles)
+        mol_img = Chem.Draw.MolToImage(mol)
+        st.image(mol_img) 
+
+    with col5:
+        smiles = 'CC(=O)Nc1ccc(O[C@@H]2O[C@H](C(=O)O)[C@@H](O)[C@H](O)[C@H]2O)cc1'
+        mol = Chem.MolFromSmiles(smiles)
+        mol_img = Chem.Draw.MolToImage(mol)
+        st.image(mol_img) 
+        
 
 def display_protein():
     """
@@ -203,7 +282,8 @@ def display_protein():
 
 page_names_to_func = {
     'About': about_page,
-    'Display DTI': display_dti
+    'Display DTI': display_dti,
+    'Retrieve Top-k': retrieval
 }
 
 selected_page = st.sidebar.selectbox('Choose function', page_names_to_func.keys())