Spaces:
Sleeping
Sleeping
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
) |