app.py

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