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