File size: 15,412 Bytes
5ebcc54
 
 
 
 
 
825b8a1
5ebcc54
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
825b8a1
 
5ebcc54
3e6c2c6
5ebcc54
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
import enum
from turtle import onclick
import streamlit as st
import numpy as np
import base64
from io import BytesIO
import requests
from multilingual_clip import pt_multilingual_clip
from transformers import CLIPTokenizerFast, AutoTokenizer
import torch
import logging
from os import environ
environ['TOKENIZERS_PARALLELISM'] = 'true'

from myscaledb import Client

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'}

@st.experimental_singleton(show_spinner=False)
def init_clip():
    """ Initialize CLIP Model

    Returns:
        Tokenizer: CLIPTokenizerFast (which convert words into embeddings)
    """
    clip = pt_multilingual_clip.MultilingualCLIP.from_pretrained(MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    return tokenizer, clip

@st.experimental_singleton(show_spinner=False)
def init_db():
    """ Initialize the Database Connection

    Returns:
        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"])
    resp = requests.get('https://www.google.com/')
    assert resp.status_code == 200
    # 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

    Args:
        xq (numpy.ndarray or list of floats): Query vector
        top_k (int, optional): Number of matched vectors. Defaults to 10.

    Returns:
        matches: list of Records object. Keys referrring to selected columns
    """
    attempt = 0
    xq = xq / np.linalg.norm(xq)
    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])
                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})")
            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
            real_xc = st.session_state.index.fetch(f"SELECT id, url, vector,\
                        distance('topK={top_k}')(vector, {xq_s}) AS dist\
                        FROM {DB_NAME}")
            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]}')
            matches = xc
            break
        except Exception as e:
            # force reload if we have trouble on connections or something else
            logging.warning(str(e))
            _, st.session_state.index = init_db()
            attempt += 1
            matches = []
    if len(matches) == 0:
        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
        self.model = torch.nn.Linear(DIMS, 1, bias=False)
        # convert initial query `xq` to tensor parameter to init weights
        init_weight = torch.Tensor(xq).reshape(1, -1)
        self.model.weight = torch.nn.Parameter(init_weight)
        # 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)
        y = torch.Tensor(y).reshape(-1, 1)
        for i in range(iters):
            # zero gradients
            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)
            # forward pass
            out = self.model(X)
            # compute loss
            loss = self.loss(out, y)
            # backward pass
            loss.backward()
            # update weights
            self.optimizer.step()
    
    def get_weights(self):
        xq = self.model.weight.detach().numpy()[0].tolist()
        return xq

def prompt2vec(prompt: str):
    """ Convert prompt into a computational vector

    Args:
        prompt (str): Text to be tokenized

    Returns:
        xq: vector from the tokenizer, representing the original prompt
    """
    # inputs = tokenizer(prompt, return_tensors='pt')
    # out = clip.get_text_features(**inputs)
    out = clip.forward(prompt, tokenizer)
    xq = out.squeeze(0).cpu().detach().numpy().tolist()
    return xq

def pil_to_bytes(img):
    """ Convert a Pillow image into base64

    Args:
        img (PIL.Image): Pillow (PIL) Image

    Returns:
        img_bin: image in base64 format
    """
    with BytesIO() as buf:
        img.save(buf, format='jpeg')
        img_bin = buf.getvalue()
        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)
    return f'<img id="img{i}" src="{url}" width="200px;" style="margin:50px 50px"><b>Relevance: {conf}</b>'

def get_top_k(xq, top_k=9):
    """ wrapper function for query

    Args:
        xq (numpy.ndarray or list of floats): Query vector
        top_k (int, optional): Number of returned vectors. Defaults to 9.

    Returns:
        matches: See `query()`
    """
    matches = query(
        xq, top_k=top_k
    )
    return matches

def tune(X, y, iters=2):
    """ Train the Zero-shot Classifier

    Args:
        X (numpy.ndarray): Input vectors (retreived vectors)
        y (list of floats or numpy.ndarray): Scores given by user
        iters (int, optional): iterations of updates to be run
    """
    # train the classifier
    st.session_state.clf.fit(X, y, iters=iters)
    # extract new vector
    st.session_state.xq = st.session_state.clf.get_weights()


def refresh_index():
    """ Clean the session
    """
    del st.session_state["meta"]
    st.session_state.meta = {}
    st.session_state.query_num = 0
    logging.info(f"Refresh for '{st.session_state.meta}'")
    init_db.clear()
    # refresh session states
    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
    scores = {}
    states = [
        st.session_state[f"input{i}"] for i in range(len(matches))
    ]
    for i, match in enumerate(matches):
        scores[match['id']] = float(states[i])
    # reset states to 1.0
    for i in range(len(matches)):
        st.session_state[f"input{i}"] = 1.0
    # get training data and labels
    X = list([match['vector'] for match in matches])
    y = [v for v in list(scores.values())]
    tune(X, y, iters=int(st.session_state.iters))
    # update record metadata after training
    for match in matches:
        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"
  href="https://fonts.googleapis.com/css?family=Roboto:300,400,500,700&display=swap"
/>
""", unsafe_allow_html=True)

messages = [
    f"""
    Find most relevant examples from a large visual dataset by combining text query and few-shot learning.
    """,
    f"""
    Then then you can adjust the weight on each image. Those weights should **represent how much it 
    can meet your preference**. You can either choose the images that match your prompt or change 
    your mind.

    You might notice that there is a iteration slide bar on the top of all retrieved images. This will
    control the speed of changes on vectors. More **iterations** will change the vector faster while 
    lower values on **iterations** will make the retrieval smoother.
    """,
    f"""
    This example will manage to train a classifier to distinguish between samples you want and samples 
    you don't want. By initializing the weight from prompt, you can get a good enough classifier to cluster
    images you want to search. If you think the result is not as perfect as you expected, you can also 
    supervise the classifer with **Relevance** annotation. If you cannot see any difference in Top-K 
    Retrieved results, try to enlarge **Number of Iteration**
    """,
    # TODO @ fangruil: fill the link with our tech blog
    f"""
    The app uses the [MyScale](http://mqdb.page.moqi.ai/mqdb-docs/) to store and query images 
    using vector search. All images are sourced from the 
    [Unsplash Lite dataset](https://unsplash-datasets.s3.amazonaws.com/lite/latest/unsplash-research-dataset-lite-latest.zip) 
    and encoded using [OpenAI's CLIP](https://huggingface.co/openai/clip-vit-base-patch32). We explain how
    it all works [here]().
    """
]

with st.spinner("Connecting DB..."):
    st.session_state.meta, st.session_state.index = init_db()

with st.spinner("Loading Models..."):
    # Initialize CLIP model
    if 'xq' not in st.session_state:
        tokenizer, clip = init_clip()
        st.session_state.query_num = 0

if 'xq' not in st.session_state:
    # If it's a fresh start
    if st.session_state.query_num < len(messages):
        msg = messages[st.session_state.query_num]
    else:
        msg = messages[-1]
        
    # Basic Layout
    
    with st.container():
        st.title("Visual Dataset Explorer")
        start = [st.empty(), st.empty(), st.empty(), st.empty(), st.empty()]
        start[0].info(msg)
        prompt = start[1].text_input("Prompt:", value="", placeholder="Examples: a photo of white dogs, cats in the snow, a house by the lake")
        start[2].markdown(
            '<p style="color:gray;"> Don\'t know what to search? Try <b>Random</b>!</p>',
            unsafe_allow_html=True)
        with start[3]:
            col = st.columns(8)
            prompt_xq = col[6].button("Prompt", disabled=len(prompt) == 0)
            random_xq = col[7].button("Random", disabled=len(prompt) != 0)
    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)
        st.session_state.xq = xq
        st.session_state.orig_xq = xq
        _ = [elem.empty() for elem in start]

if 'xq' in st.session_state:
    # If it is not a fresh start
    if st.session_state.query_num+1 < len(messages):
        msg = messages[st.session_state.query_num+1]
    else:
        msg = messages[-1]
    # 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)
    with st.container():
        with st.sidebar:
            with st.container():
                st.header("Top K Nearest in Database")
                for i, k in enumerate(st.session_state.top_k):
                    url = k["url"]
                    url += "?q=75&fm=jpg&w=200&fit=max"
                    if k["id"] not in BAD_IDS:
                        disabled = False
                    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)
                
        # 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])
            col[0].form_submit_button("Train!", on_click=submit)
            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):
                # find good url
                url = match["url"]
                url += "?q=75&fm=jpg&w=200&fit=max"
                if match["id"] not in BAD_IDS:
                    disabled = False
                else:
                    disable = True
                # the card shows an image and a checkbox
                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(
                    "Relevance",
                    min_value=0.0,
                    max_value=1.0,
                    value=1.0,
                    step=0.05,
                    key=f"input{i}",
                    disabled=disabled
                )