Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import esm
 import requests
 import matplotlib.pyplot as plt
-from myscaledb import Client
+from clickhouse_connect import get_client
 import random
 from collections import Counter
 from tqdm import tqdm
@@ -47,8 +47,6 @@ def init_db():
         host=r['host'], port=r['port'], user=st.secrets["USER"], password=st.secrets["PASSWD"],
         interface=r['http_pre'],
     )
-    # We can check if the connection is alive
-    assert client.is_alive()
     meta_field = {}
     return meta_field, Client
 
@@ -126,10 +124,8 @@ def esm_search(model, sequnce, batch_converter,top_k=5):
 
     token_list = token_representations.tolist()[0][0][0]
     
-    client = Client(
-        url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"])
-    
-    result = client.fetch("SELECT seq, distance(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer_768 ORDER BY dist LIMIT 500")
+    result = st.session_state.client.query("SELECT seq, distance(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer_768 ORDER BY dist LIMIT 500")
+    result = [r for r in result.named_results()]
     
     result_temp_seq = []
     
@@ -167,10 +163,10 @@ def KNN_search(sequence):
     token_representations = results["representations"][33]
     token_list = token_representations.tolist()[0][0] 
     print(token_list)
-    client = Client(
-        url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"])
     
-    result = client.fetch("SELECT activity, distance('topK=10')(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer")
+    result = st.session_state.client.query("SELECT activity, distance(representations, " + str(token_list) + ')'+ "as dist FROM default.esm_protein_indexer ORDER BY dist LIMIT 10")
+    result = [r for r in result.named_results()]
+    
     result_temp_activity = []
     for i in result:
         # print(result_temp_seq)
@@ -286,7 +282,7 @@ def init_random_query():
 
 
 with st.spinner("Connecting DB..."):
-    st.session_state.meta, client = init_db()
+    st.session_state.meta, st.session_state.client = init_db()
 
 with st.spinner("Loading Models..."):
     # Initialize SAGE model