added app and requirements

app.py

@@ -0,0 +1,213 @@
+import streamlit as st
+from datasets import load_dataset
+import numpy as np
+import pinecone
+import base64
+from io import BytesIO
+from transformers import CLIPTokenizerFast, CLIPModel
+import torch
+from typing import List
+
+PINECONE_API_KEY = st.secrets["PINECONE_API_KEY"]  # app.pinecone.io
+INDEX = "imagenet-query-trainer-clip"
+MODEL_ID = "openai/clip-vit-base-patch32"
+DIMS = 512
+
+@st.experimental_singleton(show_spinner=False)
+def init_dataset():
+    return load_dataset(
+        'frgfm/imagenette',
+        'full_size',
+        split='train',
+        ignore_verifications=False  # set to True if seeing splits Error
+    )
+
+@st.experimental_singleton(show_spinner=False)
+def init_clip():
+    tokenizer = CLIPTokenizerFast.from_pretrained(MODEL_ID)
+    clip = CLIPModel.from_pretrained(MODEL_ID)
+    return tokenizer, clip
+
+@st.experimental_singleton(show_spinner=False)
+def init_db():
+    pinecone.init(
+        api_key="2b2e0e96-7a31-4e87-b924-5e182ed26b03",
+        environment="us-west1-gcp"
+    )
+    return pinecone.Index(INDEX)
+
+@st.experimental_singleton(show_spinner=False)
+def init_random_query():
+    xq = np.random.rand(DIMS)
+    return xq, xq.copy()
+
+class Classifier:
+    def __init__(self, xq: list):
+        # initialize model with DIMS input size and 1 output
+        self.model = torch.nn.Linear(DIMS, 1)
+        # 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.2)
+    
+    def fit(self, X: list, y: list, iters: int = 20):
+        # 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()
+            # 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):
+    inputs = tokenizer(prompt, return_tensors='pt')
+    out = clip.get_text_features(**inputs)
+    xq = out.squeeze(0).cpu().detach().numpy().tolist()
+    return xq
+
+def pil_to_bytes(img):
+    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):
+    img = imagenet[int(i)]['image']
+    img_bin = pil_to_bytes(img)
+    return f'<img id="img{i}" src="data:image/jpeg;base64,{img_bin}" width="200px;">'
+
+def get_top_k(xq, top_k=10):
+    xc = index.query(
+        xq,
+        top_k=top_k,
+        include_values=True,
+        filter={"seen": 0}
+    )
+    matches = {match['id']: match['values'] for match in xc['matches']}
+    return matches
+
+def tune(matches, inputs):
+    positive_idx = [idx for idx, val in inputs.items() if val == 1]
+    negatives = [match for match in matches.items() if match[0] not in positive_idx]
+    negative_idx = [match[0] for match in negatives]
+    negative_vectors = [match[1] for match in negatives]
+    positive_vectors = [match[1] for match in matches.items() if match[0] in positive_idx]
+    # prep training data
+    y = [1] * len(positive_idx) + [0] * len(negative_idx)
+    X = positive_vectors + negative_vectors
+    # train the classifier
+    st.session_state.clf.fit(X, y)
+    # extract new vector
+    st.session_state.xq = st.session_state.clf.get_weights()
+    # update one record at a time
+    for i in positive_idx + negative_idx:
+        index.update(str(i), set_metadata={"seen": 1})
+    # return
+    #return clf, xq
+
+def refresh_index():
+    xq = st.session_state.xq
+    if type(xq) is not list:
+        xq = xq.tolist()
+    while True:
+        xc = index.query(xq, top_k=100, filter={"seen": 1})
+        idx = [match['id'] for match in xc['matches']]
+        if len(idx) == 0: break
+        for i in idx:
+            index.update(str(i), set_metadata={"seen": 0})
+    # refresh session states
+    del st.session_state.clf, st.session_state.xq, st.session_state.show_images
+
+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():
+    matches = st.session_state.matches
+    inputs = {}
+    states = [
+        st.session_state[f"input{i}"] for i in range(len(matches))
+    ]
+    for i, idx in enumerate(matches.keys()):
+        inputs[idx] = int(states[i])
+        states[i] = False
+    tune(matches, inputs)
+    #st.session_state.show_images = False
+
+def set_tuner_true():
+    st.session_state.tuner = True
+
+st.markdown("""
+<link
+  rel="stylesheet"
+  href="https://fonts.googleapis.com/css?family=Roboto:300,400,500,700&display=swap"
+/>
+""", unsafe_allow_html=True)
+
+with st.spinner("Initializing everything..."):
+    imagenet = init_dataset()
+    index = init_db()
+    if 'xq' not in st.session_state:
+        tokenizer, clip = init_clip()
+    if 'show_images' not in st.session_state:
+        st.session_state.show_images = False
+        st.session_state.tuner = False
+
+if 'xq' not in st.session_state:
+    prompt = st.text_input("Prompt:", value="")
+    random_xq = st.button("Random")
+    prompt_xq = st.button("Prompt", disabled=len(prompt) == 0)
+    if random_xq:
+        xq, orig_xq = init_random_query()
+        st.session_state.xq = xq
+        st.session_state.orig_xq = orig_xq
+        st.session_state.show_images = True
+    elif prompt_xq:
+        xq = prompt2vec(prompt)
+        st.session_state.xq = xq
+        st.session_state.orig_xq = xq
+        st.session_state.show_images = True
+
+else:
+    # initialize classifier
+    if 'clf' not in st.session_state:
+        st.session_state.clf = Classifier(st.session_state.xq)
+
+    new_results = st.button("Search", disabled=st.session_state.show_images)
+    if new_results:
+        st.session_state.show_images = True
+
+    refresh = st.button("Refresh")
+    if refresh:
+        # we use this to remove filters in index, refresh models etc
+        refresh_index()
+    elif st.session_state.show_images:
+        # if we want to display images we end up here
+        st.markdown(f"Distance travelled: *{round(calc_dist(), 4)}*")
+        # first retrieve images from pinecone
+        st.session_state.matches = get_top_k(st.session_state.xq, top_k=10)
+        # once retrieved, display them alongside checkboxes in a form
+        with st.form("my_form", clear_on_submit=True):
+            # we have three columns in the form
+            cols = st.columns(3)
+            for i, idx in enumerate(st.session_state.matches.keys()):
+                # the card shows an image and a checkbox
+                cols[i%3].markdown(card(idx), unsafe_allow_html=True)
+                # we access the values of the checkbox via st.session_state[f"input{i}"]
+                cols[i%3].checkbox("Relevant", key=f"input{i}")
+            st.form_submit_button("Tune", on_click=submit)

requirements.txt

@@ -0,0 +1,4 @@
+datasets
+pinecone-client
+numpy
+transformers