relation-query

.gitignore

@@ -1,2 +0,0 @@
-.streamlit/
-__pycache__

app.py

@@ -9,8 +9,7 @@ import logging
 from os import environ
 from transformers import OwlViTProcessor, OwlViTForObjectDetection
 from bot import Bot, Message
-from parse import parse
-from clickhouse_connect import get_client
+from myscaledb import Client
 from classifier import Classifier, prompt2vec, tune, SplitLayer
 from query_model import simple_query, topk_obj_query, rev_query
 from card_model import card, obj_card, style
@@ -63,11 +62,11 @@ def init_db():
         client:     Database connection object
     """
     meta = []
-    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"],
-        interface=r['http_pre'],
+    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()
     return meta, client
 
 
@@ -118,7 +117,7 @@ def query(xq, exclude_list=None):
                 IMG_DB_NAME,
                 OBJ_DB_NAME,
                 exclude_list=exclude_list,
-                topk=10,
+                topk=5000,
             )
             img_ids = [r["img_id"] for r in matches]
             if "topk_img_id" not in st.session_state:
@@ -141,7 +140,7 @@ def query(xq, exclude_list=None):
                 IMG_DB_NAME,
                 OBJ_DB_NAME,
                 thresh=-1,
-                topk=10,
+                topk=5000,
             )
             status_bar[0].write("Retrieving Non-TopK in Another TopK Images...")
             pbar.progress(75)

classifier.py

@@ -95,8 +95,8 @@ class Classifier:
             grad = []
             # Normalize the weight before inference
             # This will constrain the gradient or you will have an explosion on query vector
-            self.weight /= torch.norm(
-                self.weight, p=2, dim=-1, keepdim=True
+            self.weight.data /= torch.norm(
+                self.weight.data, p=2, dim=-1, keepdim=True
             )
             for n in range(self.num_class):
                 # select all training sample and create labels
@@ -109,25 +109,22 @@ class Classifier:
                 # To simplify the query, we separated 
                 # the calculation into class numbers
                 grad_q_str = f"""
-                    SELECT avgForEachArray(arrayMap((x,y,gt)->arrayMap(i->i*(y-gt), x), X, Y, GT)) AS grad
+                    SELECT sumForEachArray(arrayMap((x,y,gt)->arrayMap(i->i*(y-gt), x), X, Y, GT)) AS grad
                     FROM (
                         SELECT groupArray(arrayPopBack(prelogit)) AS X, 
                             groupArray(1/(1+exp(-arraySum(arrayMap((x,y)->x*y, prelogit, {xq_s[n]}))))) AS Y, {labels} AS GT
                         FROM {self.obj_db} WHERE obj_id IN {objs})"""
-                grad_ = [r['grad'] for r in self.client.query(grad_q_str).named_results()][0]
-                grad.append(torch.as_tensor(grad_))
+                grad.append(torch.as_tensor(self.client.fetch(grad_q_str)[0]['grad']))
             # update weights
             grad = torch.stack(grad, dim=0)
-            self.weight -= 0.01 * grad
-            self.weight /= torch.norm(
-                self.weight, p=2, dim=-1, keepdim=True
-            )
+            self.weight -= 0.1 * grad
 
     def get_weights(self):
         xq = self.weight.detach().numpy()
         return xq
 
 
+
 class SplitLayer(torch.nn.Module):
     def forward(self, x):
         return torch.split(x, 1, dim=-1)

query_model.py

@@ -19,11 +19,11 @@ def topk_obj_query(client, xq, IMG_DB_NAME, OBJ_DB_NAME,
             FROM {OBJ_DB_NAME}
             JOIN {IMG_DB_NAME}
             ON {IMG_DB_NAME}.img_id = {OBJ_DB_NAME}.img_id
-            WHERE obj_id IN (
+            PREWHERE obj_id IN (
                 SELECT obj_id FROM (
-                    SELECT obj_id, distance(prelogit, {_xq}) AS dist FROM {OBJ_DB_NAME}
-                    ORDER BY dist DESC LIMIT 5000
-                ) {_cond} LIMIT {topk}
+                    SELECT obj_id, distance('topK={topk}', 'nprobe=32')(prelogit, {_xq}) AS dist FROM {OBJ_DB_NAME}
+                    ORDER BY dist DESC
+                ) {_cond} LIMIT 10
             )
             """)
     _subq_str = ' UNION ALL '.join(_subq_str)
@@ -38,7 +38,7 @@ def topk_obj_query(client, xq, IMG_DB_NAME, OBJ_DB_NAME,
                     ({_subq_str})
             GROUP BY img_id, img_url, img_w, img_h ORDER BY img_score DESC
     """
-    xc = [{k: v for k, v in r.items()} for r in client.query(q_str).named_results()]
+    xc = client.fetch(q_str)
     return xc
 
 
@@ -74,7 +74,7 @@ def rev_query(client, xq, img_ids, IMG_DB_NAME, OBJ_DB_NAME, thresh=0.08):
                 ({_subq_str})
             GROUP BY img_id, img_url, img_w, img_h ORDER BY img_score DESC
     """
-    xc = [{k: v for k, v in r.items()} for r in client.query(q_str).named_results()]
+    xc = client.fetch(q_str)
     return xc
 
 
@@ -88,11 +88,11 @@ def simple_query(client, xq, IMG_DB_NAME, OBJ_DB_NAME, thresh=0.08, topk=10):
         subq_str.append(
             f"""
             SELECT {OBJ_DB_NAME}.img_id AS img_id, img_url, img_w, img_h, prelogit,
-                obj_id, box_cx, box_cy, box_w, box_h, {_l} AS l,  distance(prelogit, {_xq}) AS dist
+                obj_id, box_cx, box_cy, box_w, box_h, {_l} AS l,  distance('topK={topk}', 'nprobe=32')(prelogit, {_xq}) AS dist
             FROM {OBJ_DB_NAME}
             JOIN {IMG_DB_NAME} 
             ON {IMG_DB_NAME}.img_id = {OBJ_DB_NAME}.img_id
-            {_thresh} ORDER BY dist DESC LIMIT {topk}
+            {_thresh} LIMIT 10
             """)
     subq_str = " UNION ALL ".join(subq_str)
     q_str = f"""
@@ -104,5 +104,5 @@ def simple_query(client, xq, IMG_DB_NAME, OBJ_DB_NAME, thresh=0.08, topk=10):
                 )
                 GROUP BY l
                 """
-    xc = [{k: v for k, v in r.items()} for r in client.query(q_str).named_results()]
-    return xc
+    res = client.fetch(q_str)
+    return res

requirements.txt

@@ -1,9 +1,7 @@
 transformers
 tqdm
-clickhouse-connect
-parse
+myscaledb-client==1.1.7
 streamlit
-altair < 5
 numpy
 torch
 onnx