Update app.py

app.py

@@ -8,6 +8,7 @@ from collections import Counter
 from tqdm import tqdm
 from statistics import mean
 
+import torch
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -16,6 +17,8 @@ from stmol import *
 import py3Dmol
 # from streamlit_3Dmol import component_3dmol
 
+import esm
+
 import scipy
 from sklearn.model_selection import GridSearchCV, train_test_split
 from sklearn.decomposition import PCA
@@ -115,14 +118,17 @@ def esm_search(model, sequnce, batch_converter,top_k=5):
     client = Client(
         url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"])
     
-    result = client.fetch("SELECT activity, distance('topK=5')(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer_768")
+    result = client.fetch("SELECT seq, distance('topK=500')(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer_768")
+    
     result_temp_seq = []
+    
     for i in result:
-        # print(result_temp_seq)
-        result_temp_coords = i['coords']
+        # result_temp_coords = i['seq']
         result_temp_seq.append(i['seq'])
+    
+    result_temp_seq = list(set(result_temp_seq))
 
-    return result_temp_coords, result_temp_seq
+    return result_temp_seq
 
 def KNN_search(sequence):
     model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
@@ -307,7 +313,7 @@ if 'xq' not in st.session_state:
             
             if sequence:
                 st.write('you have entered: ', sequence)
-                result_temp_coords, result_temp_seq = esm_search(model, sequence, esm_search,top_k=5)
+                result_temp_seq = esm_search(model, sequence, esm_search,top_k=5)
                 st.text('search result: ')
                 # tab1, tab2, tab3, tab4, = st.tabs(["Cat", "Dog", "Owl"])
                 if st.button(result_temp_seq[0]):
@@ -380,11 +386,13 @@ else:
             
             if sequence:
                 st.write('you have entered: ', sequence)
-                result_temp_coords, result_temp_seq = esm_search(st.session_state['xq'], sequence, st.session_state['batch'] ,top_k=1)
+                result_temp_seq = esm_search(st.session_state['xq'], sequence, st.session_state['batch'] ,top_k=10)
                 st.text('search result (top 5): ')
                 # tab1, tab2, tab3, tab4, = st.tabs(["Cat", "Dog", "Owl"])
-                option2 = st.selectbox('top5 sequence', (result_temp_seq[0],result_temp_seq[1],result_temp_seq[2],result_temp_seq[3],result_temp_seq[4]))
-                if option2 == result_temp_seq[0]:
+                tab1, tab2, tab3 , tab4, tab5 = st.tabs(['1','2','3','4','5'])
+                # option2 = st.radio('top5 sequence', (result_temp_seq[0],result_temp_seq[1],result_temp_seq[2],result_temp_seq[3],result_temp_seq[4]))
+
+                with tab1:
                     st.write(result_temp_seq[0])
                     import random
                     # print(random.randint(0,9))
@@ -393,8 +401,9 @@ else:
                     protein = prot_str[random.randint(14,18)]
                     xyzview = py3Dmol.view(query='pdb:'+protein)
                     xyzview.setStyle({'stick':{'color':'spectrum'}})
-                    start[3] = showmol(xyzview, height = 500,width=800)
+                    showmol(xyzview, height = 500,width=800)
                      # st.write(result_temp_seq[4])
+                with tab2:
                     import random
                     # print(random.randint(0,9))
                     st.write(result_temp_seq[1])
@@ -403,28 +412,31 @@ else:
                     protein = prot_str[random.randint(0,4)]
                     xyzview = py3Dmol.view(query='pdb:'+protein)
                     xyzview.setStyle({'stick':{'color':'spectrum'}})
-                    start[4] = showmol(xyzview, height = 500,width=800)
+                    showmol(xyzview, height = 500,width=800)
+                with tab3:
                     st.write(result_temp_seq[2])
                     prot_str=['1A2C','1BML','1D5M','1D5X','1D5Z','1D6E','1DEE','1E9F','1FC2','1FCC','1G4U','1GZS','1HE1','1HEZ','1HQR','1HXY','1IBX','1JBU','1JWM','1JWS']
                     # protein=st.selectbox('select protein',prot_list)
                     protein = prot_str[random.randint(4,8)]
                     xyzview = py3Dmol.view(query='pdb:'+protein)
                     xyzview.setStyle({'stick':{'color':'spectrum'}})
-                    start[5] = showmol(xyzview, height = 500,width=800)
+                    showmol(xyzview, height = 500,width=800)
+                with tab4:
                     st.write(result_temp_seq[3])
                     prot_str=['1A2C','1BML','1D5M','1D5X','1D5Z','1D6E','1DEE','1E9F','1FC2','1FCC','1G4U','1GZS','1HE1','1HEZ','1HQR','1HXY','1IBX','1JBU','1JWM','1JWS']
                     # protein=st.selectbox('select protein',prot_list)
                     protein = prot_str[random.randint(4,8)]
                     xyzview = py3Dmol.view(query='pdb:'+protein)
                     xyzview.setStyle({'stick':{'color':'spectrum'}})
-                    start[6] = showmol(xyzview, height = 500,width=800)
+                    showmol(xyzview, height = 500,width=800)
+                with tab5:
                     st.write(result_temp_seq[4])
                     prot_str=['1A2C','1BML','1D5M','1D5X','1D5Z','1D6E','1DEE','1E9F','1FC2','1FCC','1G4U','1GZS','1HE1','1HEZ','1HQR','1HXY','1IBX','1JBU','1JWM','1JWS']
                     # protein=st.selectbox('select protein',prot_list)
                     protein = prot_str[random.randint(4,8)]
                     xyzview = py3Dmol.view(query='pdb:'+protein)
                     xyzview.setStyle({'stick':{'color':'spectrum'}})
-                    start[7] = showmol(xyzview, height = 500,width=800)
+                    showmol(xyzview, height = 500,width=800)
                 
 
         elif option == 'activity prediction':