add features and datasets

.gitignore

@@ -0,0 +1 @@
+.streamlit

app.py

@@ -7,15 +7,22 @@ from transformers import CLIPTokenizerFast, AutoTokenizer
 import torch
 import logging
 from os import environ
+from myscaledb import Client
 environ['TOKENIZERS_PARALLELISM'] = 'true'
 
-from myscaledb import Client
+
+db_name_map = {
+    "Unsplash Photos 25K": "mqdb_demo.unsplash_25k_clip_indexer",
+    "RSICD: Remote Sensing Images 11K": "mqdb_demo.rsicd_clip_b_32",
+}
 
 DB_NAME = "mqdb_demo.unsplash_25k_clip_indexer"
 MODEL_ID = 'M-CLIP/XLM-Roberta-Large-Vit-B-32'
 DIMS = 512
 # Ignore some bad links (broken in the dataset already)
-BAD_IDS = {'9_9hzZVjV8s', 'RDs0THr4lGs', 'vigsqYux_-8', 'rsJtMXn3p_c', 'AcG-unN00gw', 'r1R_0ZNUcx0'}
+BAD_IDS = {'9_9hzZVjV8s', 'RDs0THr4lGs', 'vigsqYux_-8',
+           'rsJtMXn3p_c', 'AcG-unN00gw', 'r1R_0ZNUcx0'}
+
 
 @st.experimental_singleton(show_spinner=False)
 def init_clip():
@@ -28,6 +35,7 @@ def init_clip():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
     return tokenizer, clip
 
+
 @st.experimental_singleton(show_spinner=False)
 def init_db():
     """ Initialize the Database Connection
@@ -36,17 +44,20 @@ 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"])
+    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()
     meta_field = {}
     return meta_field, client
 
+
 @st.experimental_singleton(show_spinner=False)
 def init_query_num():
     print("init query_num")
     return 0
 
+
 def query(xq, top_k=10):
     """ Query TopK matched w.r.t a given vector
 
@@ -62,30 +73,29 @@ def query(xq, top_k=10):
     while attempt < 3:
         try:
             xq_s = f"[{', '.join([str(float(fnum)) for fnum in list(xq)])}]"
-            
+
             print('Excluded pre:', st.session_state.meta)
             if len(st.session_state.meta) > 0:
-                exclude_list = ','.join([f'\'{i}\'' for i, v in st.session_state.meta.items() if v >= 1])
+                exclude_list = ','.join(
+                    [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,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {DB_NAME} PREWHERE id NOT IN ({exclude_list})")
+                        FROM {db_name_map[st.session_state.db_name_ref]} \
+                        PREWHERE id NOT IN ({exclude_list})")
             else:
                 xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {DB_NAME}")
-            # real_xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
-            #             1 - arraySum(arrayMap((x, y) -> x * y, {xq_s}, vector)) AS dist\
-            #             FROM {DB_NAME} ORDER BY dist LIMIT {top_k}")
-            # FIXME: This is causing freezing on DB
+                        FROM {db_name_map[st.session_state.db_name_ref]}")
             real_xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                         distance('topK={top_k}')(vector, {xq_s}) AS dist\
-                        FROM {DB_NAME}")
+                        FROM {db_name_map[st.session_state.db_name_ref]}")
             top_k = 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(f'{len(xc)} records returned, {[_i["id"] for _i in 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(
+                f'{len(xc)} records returned, {[_i["id"] for _i in xc]}')
             matches = xc
             break
         except Exception as e:
@@ -98,20 +108,23 @@ def query(xq, top_k=10):
         logging.error(f"No matches found for '{DB_NAME}'")
     return matches, top_k
 
+
 @st.experimental_singleton(show_spinner=False)
 def init_random_query():
     xq = np.random.rand(DIMS).tolist()
     return xq, xq.copy()
 
+
 class Classifier:
     """ Zero-shot Classifier
     This Classifier provides proxy regarding to the user's reaction to the probed images.
     The proxy will replace the original query vector generated by prompted vector and finally
     give the user a satisfying retrieval result.
-    
+
     This can be commonly seen in a recommendation system. The classifier will recommend more 
     precise result as it accumulating user's activity.
     """
+
     def __init__(self, xq: list):
         # initialize model with DIMS input size and 1 output
         # note that the bias is ignored, as we only focus on the inner product result
@@ -122,7 +135,7 @@ class Classifier:
         # init loss and optimizer
         self.loss = torch.nn.BCEWithLogitsLoss()
         self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
-    
+
     def fit(self, X: list, y: list, iters: int = 5):
         # convert X and y to tensor
         X = torch.Tensor(X)
@@ -132,7 +145,8 @@ class Classifier:
             self.optimizer.zero_grad()
             # Normalize the weight before inference
             # This will constrain the gradient or you will have an explosion on query vector
-            self.model.weight.data = self.model.weight.data / torch.norm(self.model.weight.data, p=2, dim=-1)
+            self.model.weight.data = self.model.weight.data / \
+                torch.norm(self.model.weight.data, p=2, dim=-1)
             # forward pass
             out = self.model(X)
             # compute loss
@@ -141,11 +155,17 @@ class Classifier:
             loss.backward()
             # update weights
             self.optimizer.step()
-    
+
     def get_weights(self):
         xq = self.model.weight.detach().numpy()[0].tolist()
         return xq
 
+
+class NormalizingLayer(torch.nn.Module):
+    def forward(self, x):
+        return x / torch.norm(x, dim=-1, keepdim=True)
+
+
 def prompt2vec(prompt: str):
     """ Convert prompt into a computational vector
 
@@ -161,6 +181,7 @@ def prompt2vec(prompt: str):
     xq = out.squeeze(0).cpu().detach().numpy().tolist()
     return xq
 
+
 def pil_to_bytes(img):
     """ Convert a Pillow image into base64
 
@@ -176,13 +197,16 @@ def pil_to_bytes(img):
         img_bin = base64.b64encode(img_bin).decode('utf-8')
     return img_bin
 
+
 def card(i, url):
     return f'<img id="img{i}" src="{url}" width="200px;">'
 
+
 def card_with_conf(i, conf, url):
-    conf = "%.4f"%(conf)
+    conf = "%.4f" % (conf)
     return f'<img id="img{i}" src="{url}" width="200px;" style="margin:50px 50px"><div><p><b>Relevance: {conf}</b></p></div>'
 
+
 def get_top_k(xq, top_k=9):
     """ wrapper function for query
 
@@ -198,6 +222,7 @@ def get_top_k(xq, top_k=9):
     )
     return matches
 
+
 def tune(X, y, iters=2):
     """ Train the Zero-shot Classifier
 
@@ -206,6 +231,7 @@ def tune(X, y, iters=2):
         y (list of floats or numpy.ndarray): Scores given by user
         iters (int, optional): iterations of updates to be run
     """
+    assert len(X) == len(y)
     # train the classifier
     st.session_state.clf.fit(X, y, iters=iters)
     # extract new vector
@@ -224,17 +250,19 @@ def refresh_index():
     st.session_state.meta, st.session_state.index = init_db()
     del st.session_state.clf, st.session_state.xq
 
+
 def calc_dist():
     xq = np.array(st.session_state.xq)
     orig_xq = np.array(st.session_state.orig_xq)
     return np.linalg.norm(xq - orig_xq)
 
+
 def submit():
     """ Tune the model w.r.t given score from user.
     """
     st.session_state.query_num += 1
     matches = st.session_state.matches
-    velocity = 1 #st.session_state.velocity
+    velocity = 1  # st.session_state.velocity
     scores = {}
     states = [
         st.session_state[f"input{i}"] for i in range(len(matches))
@@ -253,9 +281,11 @@ def submit():
         st.session_state.meta[match['id']] = 1
     logging.info(f"Exclude List: {st.session_state.meta}")
 
+
 def delete_element(element):
     del element
 
+
 st.markdown("""
 <link
   rel="stylesheet"
@@ -308,32 +338,56 @@ if 'xq' not in st.session_state:
         msg = messages[st.session_state.query_num]
     else:
         msg = messages[-1]
-        
+    prompt = ''
     # Basic Layout
-    
     with st.container():
+        if 'prompt' in st.session_state:
+            del st.session_state.prompt
         st.title("Visual Dataset Explorer")
-        start = [st.empty(), st.empty(), st.empty(), st.empty(), st.empty()]
+        start = [st.empty(), st.empty(), st.empty(), st.empty(),
+                 st.empty(), st.empty(), st.empty()]
         start[0].info(msg)
-        prompt = start[1].text_input("Prompt:", value="", placeholder="Examples: white dogs, 女人举着雨伞, mouette volant au-dessus de la mer, ガラスの花瓶の花 ...")
-        start[2].markdown(
+        st.session_state.db_name_ref = start[1].selectbox(
+            "Select Database:", list(db_name_map.keys()))
+        prompt = start[2].text_input(
+            "Prompt:", value="", placeholder="Examples: playing corgi, 女人举着雨伞, mouette volant au-dessus de la mer, ガラスの花瓶の花 ...")
+        if len(prompt) > 0:
+            st.session_state.prompt = prompt
+        start[3].markdown(
             '<p style="color:gray;"> Don\'t know what to search? Try <b>Random</b>!</p>\
             <p>🌟 We also support multi-language search. Type any language you know to search! ⌨️ </p>',
             unsafe_allow_html=True)
-        with start[3]:
+        upld_model = start[5].file_uploader(
+            "Or you can upload your previous run!", type='onnx')
+        upld_btn = start[6].button(
+            "Used Loaded Weights", disabled=upld_model is None)
+        with start[4]:
             col = st.columns(8)
+            has_no_prompt = (len(prompt) == 0 and upld_model is None)
             prompt_xq = col[6].button("Prompt", disabled=len(prompt) == 0)
-            random_xq = col[7].button("Random", disabled=len(prompt) != 0)
+            random_xq = col[7].button("Random", disabled=not (
+                len(prompt) == 0 and upld_model is None))
+
     if random_xq:
         # Randomly pick a vector to query
         xq, orig_xq = init_random_query()
         st.session_state.xq = xq
         st.session_state.orig_xq = orig_xq
         _ = [elem.empty() for elem in start]
-    elif prompt_xq:
-        print(f"Input prompt is {prompt}")
-        # Tokenize the vectors
-        xq = prompt2vec(prompt)
+    elif prompt_xq or upld_btn:
+        if upld_model is not None:
+            # Import vector from a file
+            import onnx
+            from onnx import numpy_helper
+            _model = onnx.load(upld_model)
+            weights = _model.graph.initializer
+            assert len(weights) == 1
+            xq = numpy_helper.to_array(weights[0]).tolist()
+            assert len(xq) == DIMS
+        else:
+            print(f"Input prompt is {prompt}")
+            # Tokenize the vectors
+            xq = prompt2vec(prompt)
         st.session_state.xq = xq
         st.session_state.orig_xq = xq
         _ = [elem.empty() for elem in start]
@@ -347,11 +401,21 @@ if 'xq' in st.session_state:
     # initialize classifier
     if 'clf' not in st.session_state:
         st.session_state.clf = Classifier(st.session_state.xq)
-    
+
     # if we want to display images we end up here
     st.info(msg)
     # first retrieve images from pinecone
-    st.session_state.matches, st.session_state.top_k = get_top_k(st.session_state.clf.get_weights(), top_k=9)
+    st.session_state.matches, st.session_state.top_k = get_top_k(
+        st.session_state.clf.get_weights(), top_k=9)
+
+    # export the model into executable ONNX
+    st.session_state.dnld_model = BytesIO()
+    torch.onnx.export(torch.nn.Sequential(NormalizingLayer(), st.session_state.clf.model),
+                      torch.as_tensor(st.session_state.xq).reshape(1, -1),
+                      st.session_state.dnld_model,
+                      input_names=['input'],
+                      output_names=['output'])
+
     with st.container():
         with st.sidebar:
             with st.container():
@@ -364,15 +428,23 @@ if 'xq' in st.session_state:
                     else:
                         disable = True
                     dist = np.matmul(st.session_state.clf.get_weights() / np.linalg.norm(st.session_state.clf.get_weights()),
-                                    np.array(k["vector"]).T)
-                    st.markdown(card_with_conf(i, dist, url), unsafe_allow_html=True)
-                
+                                     np.array(k["vector"]).T)
+                    st.markdown(card_with_conf(i, dist, url),
+                                unsafe_allow_html=True)
+        dnld_nam = st.text_input('Download Name:',
+                                 f'{(st.session_state.prompt if "prompt" in st.session_state else (upld_model.name.split(".onnx")[0] if upld_model is not None else "model"))}.onnx',
+                                 max_chars=50)
+        dnld_btn = st.download_button('Download your classifier!',
+                                      st.session_state.dnld_model,
+                                      dnld_nam,)
         # once retrieved, display them alongside checkboxes in a form
         with st.form("batch", clear_on_submit=False):
-            st.session_state.iters = st.slider("Number of Iterations to Update", min_value=0, max_value=10, step=1, value=2)
-            col = st.columns([1,9])
+            st.session_state.iters = st.slider(
+                "Number of Iterations to Update", min_value=0, max_value=10, step=1, value=2)
+            col = st.columns([1, 9])
             col[0].form_submit_button("Train!", on_click=submit)
-            col[1].form_submit_button("Choose a new prompt", on_click=refresh_index)
+            col[1].form_submit_button(
+                "Choose a new prompt", on_click=refresh_index)
             # we have three columns in the form
             cols = st.columns(3)
             for i, match in enumerate(st.session_state.matches):
@@ -384,9 +456,9 @@ if 'xq' in st.session_state:
                 else:
                     disable = True
                 # the card shows an image and a checkbox
-                cols[i%3].markdown(card(i, url), unsafe_allow_html=True)
+                cols[i % 3].markdown(card(i, url), unsafe_allow_html=True)
                 # we access the values of the checkbox via st.session_state[f"input{i}"]
-                cols[i%3].slider(
+                cols[i % 3].slider(
                     "Relevance",
                     min_value=0.0,
                     max_value=1.0,
@@ -394,4 +466,4 @@ if 'xq' in st.session_state:
                     step=0.05,
                     key=f"input{i}",
                     disabled=disabled
-                )
+                )

requirements.txt

@@ -4,4 +4,5 @@ myscaledb-client
 streamlit
 multilingual-clip
 numpy
-torch
+torch
+onnx