switch to new ch client

app.py

@@ -7,7 +7,8 @@ from transformers import CLIPTokenizerFast, AutoTokenizer, CLIPModel
 import torch
 import logging
 from os import environ
-from myscaledb import Client
+from parse import parse
+from clickhouse_connect import get_client
 environ['TOKENIZERS_PARALLELISM'] = 'true'
 
 
@@ -36,10 +37,10 @@ def init_db():
         meta_field: Meta field that records if an image is viewed or not
         client:     Database connection object
     """
-    client = Client(
-        url=st.secrets["DB_URL"], user=st.secrets["USER"], password=st.secrets["PASSWD"])
-    # We can check if the connection is alive
-    assert client.is_alive()
+    r = parse("{http_pre}://{host}:{port}", st.secrets["DB_URL"])
+    client = get_client(
+        host=r['host'], port=r['port'], user=st.secrets["USER"], password=st.secrets["PASSWD"]
+    )
     meta_field = {}
     return meta_field, client
 
@@ -72,18 +73,18 @@ def query(xq, top_k=10):
                     [f'\'{i}\'' for i, v in st.session_state.meta.items() if v >= 1])
                 print("Excluded:", exclude_list)
                 # Using PREWHERE allows you to do column filter before vector search
-                xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
+                xc = st.session_state.index.query(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
                         FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])} \
-                        PREWHERE id NOT IN ({exclude_list})")
+                        PREWHERE id NOT IN ({exclude_list})").named_results()
             else:
-                xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
+                xc = st.session_state.index.query(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}")
-            real_xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
+                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}").named_results()
+            real_xc = st.session_state.index.query(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}")
-            top_k = real_xc
+                        FROM {db_name_map[st.session_state.db_name_ref](feat_name_map[st.session_state.feat_name])}").named_results()
+            top_k = [{k: v for k, v in r.items()} for r in real_xc]
             xc = [xi for xi in xc if xi['id'] not in st.session_state.meta or
                   st.session_state.meta[xi['id']] < 1]
             logging.info(

requirements.txt

@@ -1,6 +1,7 @@
 transformers
 tqdm
-myscaledb-client==1.1.7
+parse
+clickhouse-connect
 streamlit
 multilingual-clip
 numpy