import enum
from turtle import onclick
import streamlit as st
import numpy as np
import base64
from io import BytesIO
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"])
# 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''
def card_with_conf(i, conf, url):
conf = "%.4f"%(conf)
return f'Relevance: {conf}'
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("""
""", 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(
'
Don\'t know what to search? Try Random!
',
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
)