add the application files

Dockerfile

@@ -0,0 +1,7 @@
+FROM python:3.8-slim-buster
+COPY . /app
+WORKDIR /app
+RUN pip install -r requirements.txt
+EXPOSE 8501
+ENTRYPOINT ["streamlit","run"]
+CMD ["app.py"]

README.md

@@ -1,8 +1,8 @@
 ---
 title: Image Search
-emoji: 🔥
-colorFrom: pink
-colorTo: indigo
+emoji: 😻
+colorFrom: blue
+colorTo: gray
 sdk: streamlit
 app_file: app.py
 pinned: false

SessionState.py

@@ -0,0 +1,107 @@
+"""Hack to add per-session state to Streamlit.
+Usage
+-----
+>>> import SessionState
+>>>
+>>> session_state = SessionState.get(user_name='', favorite_color='black')
+>>> session_state.user_name
+''
+>>> session_state.user_name = 'Mary'
+>>> session_state.favorite_color
+'black'
+Since you set user_name above, next time your script runs this will be the
+result:
+>>> session_state = get(user_name='', favorite_color='black')
+>>> session_state.user_name
+'Mary'
+"""
+try:
+    import streamlit.ReportThread as ReportThread
+    from streamlit.server.Server import Server
+except Exception:
+    # Streamlit >= 0.65.0
+    import streamlit.report_thread as ReportThread
+    from streamlit.server.server import Server
+
+
+class SessionState(object):
+    def __init__(self, **kwargs):
+        """A new SessionState object.
+        Parameters
+        ----------
+        **kwargs : any
+            Default values for the session state.
+        Example
+        -------
+        >>> session_state = SessionState(user_name='', favorite_color='black')
+        >>> session_state.user_name = 'Mary'
+        ''
+        >>> session_state.favorite_color
+        'black'
+        """
+        for key, val in kwargs.items():
+            setattr(self, key, val)
+
+
+def get(**kwargs):
+    """Gets a SessionState object for the current session.
+    Creates a new object if necessary.
+    Parameters
+    ----------
+    **kwargs : any
+        Default values you want to add to the session state, if we're creating a
+        new one.
+    Example
+    -------
+    >>> session_state = get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    ''
+    >>> session_state.user_name = 'Mary'
+    >>> session_state.favorite_color
+    'black'
+    Since you set user_name above, next time your script runs this will be the
+    result:
+    >>> session_state = get(user_name='', favorite_color='black')
+    >>> session_state.user_name
+    'Mary'
+    """
+    # Hack to get the session object from Streamlit.
+
+    ctx = ReportThread.get_report_ctx()
+
+    this_session = None
+
+    current_server = Server.get_current()
+    if hasattr(current_server, '_session_infos'):
+        # Streamlit < 0.56
+        session_infos = Server.get_current()._session_infos.values()
+    else:
+        session_infos = Server.get_current()._session_info_by_id.values()
+
+    for session_info in session_infos:
+        s = session_info.session
+        if (
+            # Streamlit < 0.54.0
+            (hasattr(s, '_main_dg') and s._main_dg == ctx.main_dg)
+            or
+            # Streamlit >= 0.54.0
+            (not hasattr(s, '_main_dg') and s.enqueue == ctx.enqueue)
+            or
+            # Streamlit >= 0.65.2
+            (not hasattr(s, '_main_dg') and s._uploaded_file_mgr == ctx.uploaded_file_mgr)
+        ):
+            this_session = s
+
+    if this_session is None:
+        raise RuntimeError(
+            "Oh noes. Couldn't get your Streamlit Session object. "
+            'Are you doing something fancy with threads?')
+
+    # Got the session object! Now let's attach some state into it.
+
+    if not hasattr(this_session, '_custom_session_state'):
+        this_session._custom_session_state = SessionState(**kwargs)
+
+    return this_session._custom_session_state
+
+__all__ = ['get']

app.py

@@ -0,0 +1,83 @@
+import streamlit as st
+import SessionState
+from prompts import PROMPT_LIST
+from wit_index import WitIndex
+import random
+import time
+
+# st.set_page_config(page_title="Image Search")
+
+# vector_length = 128
+wit_index_path = f"./models/wit_faiss.idx"
+model_name = f"./models/distilbert-base-wit"
+wit_dataset_path = "./models/wit_dataset.pkl"
+
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def get_wit_index():
+    st.write("Loading the WIT index, dataset and the DistillBERT model..")
+    wit_index = WitIndex(wit_index_path, model_name, wit_dataset_path, gpu=False)
+    return wit_index
+
+@st.cache(suppress_st_warning=True)
+def process(text: str, top_k: int = 10):
+    # st.write("Cache miss: process")
+    distance, index, image_info = wit_index.search(text, top_k=top_k)
+    return distance, index, image_info
+
+
+st.title("Image Search")
+
+st.markdown(
+    """
+    This application is a demo for sentence-based image search using 
+    [WIT dataset](https://github.com/google-research-datasets/wit). We use DistillBert to encode the sentences 
+    and Facebook's Faiss to search the vector embeddings.
+    """
+)
+session_state = SessionState.get(prompt=None, prompt_box=None, text=None)
+ALL_PROMPTS = list(PROMPT_LIST.keys())+["Custom"]
+prompt = st.selectbox('Prompt', ALL_PROMPTS, index=len(ALL_PROMPTS)-1)
+# Update prompt
+if session_state.prompt is None:
+    session_state.prompt = prompt
+elif session_state.prompt is not None and (prompt != session_state.prompt):
+    session_state.prompt = prompt
+    session_state.prompt_box = None
+    session_state.text = None
+else:
+    session_state.prompt = prompt
+
+# Update prompt box
+if session_state.prompt == "Custom":
+    session_state.prompt_box = "Enter your text here"
+else:
+    if session_state.prompt is not None and session_state.prompt_box is None:
+        session_state.prompt_box = random.choice(PROMPT_LIST[session_state.prompt])
+
+session_state.text = st.text_area("Enter text", session_state.prompt_box)
+
+top_k = st.sidebar.number_input(
+    "Top k",
+    value=6,
+    min_value=1,
+    max_value=10
+)
+
+wit_index = get_wit_index()
+if st.button("Run"):
+    with st.spinner(text="Getting results..."):
+        st.subheader("Result")
+        time_start = time.time()
+        distances, index, image_info = process(text=session_state.text, top_k=int(top_k))
+        time_end = time.time()
+        time_diff = time_end-time_start
+        print(f"Search in {time_diff} seconds")
+        st.markdown(f"*Search in {time_diff:.5f} seconds*")
+        for i, distance in enumerate(distances):
+            st.image(image_info[i][0].replace("http:", "https:"), width=400)
+            st.write(f"{image_info[i][1]}. (D: {distance:.2f})")
+
+        # Reset state
+        session_state.prompt = None
+        session_state.prompt_box = None
+        session_state.text = None

prompts.py

@@ -0,0 +1,18 @@
+PROMPT_LIST = {
+    "City": [
+        "Vienna is the national capital, largest city, and one of nine states of Austria. Vienna is Austria's most populous city, with about 2 million inhabitants, and its cultural, economic, and political centre. It is the 6th-largest city by population within city limits in the European Union",
+        "Sydney is the capital city of the state of New South Wales, and the most populous city in Australia and Oceania.",
+        "Ubud is a town on the Indonesian island of Bali in Ubud District, located amongst rice paddies and steep ravines in the central foothills of the Gianyar regency. Promoted as an arts and culture centre, it has developed a large tourism industry.",
+        "Jakarta is the capital of Indonesia"
+    ],
+    "People": [
+        "Albert Einstein was a German-born theoretical physicist, widely acknowledged to be one of the greatest physicists of all time. Einstein is known for developing the theory of relativity, but he also made important contributions to the development of the theory of quantum mechanics.",
+        "Geoffrey Everest Hinton is a British-Canadian cognitive psychologist and computer scientist, most noted for his work on artificial neural networks.",
+        "Pramoedya Ananta Toer was an Indonesian author of novels, short stories, essays, polemics and histories of his homeland and its people."
+    ],
+    "Building": [
+        "Borobudur is a 7th-century Mahayana Buddhist temple in Indonesia. It is the world's largest Buddhist temple. The temple consists of nine stacked platforms, six square and three circular, topped by a central dome. It is decorated with 2,672 relief panels and 504 Buddha statues.",
+        "The Statue of Liberty is a colossal neoclassical sculpture on Liberty Island in New York Harbor within New York City, in the United States.",
+        "Machu Picchu is a 15th-century Inca citadel, located in the Eastern Cordillera of southern Peru, on a 2,430-meter mountain ridge."
+    ]
+}

requirements.txt

@@ -0,0 +1,11 @@
+numpy
+torch
+tensorflow
+transformers
+tensorflow
+mtranslate
+sentence-transformers
+datasets
+faiss-cpu
+# streamlit version 0.67.1 is needed due to issue with caching
+streamlit==0.67.1

wit_index.py

@@ -0,0 +1,40 @@
+import pickle
+
+# Used to create the dense document vectors.
+import torch
+from sentence_transformers import SentenceTransformer
+import datasets
+
+# Used to create and store the Faiss index.
+import faiss
+import numpy as np
+
+class WitIndex:
+    """
+    WitIndex is a class to search the wiki snippets from the given text. It can also return link to the
+    wiki page or the image.
+    """
+    wit_dataset = None
+
+    def __init__(self, wit_index_path: str, model_name: str, wit_dataset_path: str, gpu=True):
+        self.index = faiss.read_index(wit_index_path)
+        self.model = SentenceTransformer(model_name)
+        if WitIndex.wit_dataset is None:
+            WitIndex.wit_dataset = pickle.load(open(wit_dataset_path, "rb"))
+        print(f"Gpu: {gpu}")
+        if gpu and torch.cuda.is_available():
+            print("Cuda is available")
+            self.model = self.model.to(torch.device("cuda"))
+            res = faiss.StandardGpuResources()
+            self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
+
+    def search(self, text, top_k=6):
+        print(f"> Search: {text}")
+        embedding = self.model.encode(text, convert_to_numpy=True, show_progress_bar=False)
+        # Retrieve the k nearest neighbours
+        distance, index = self.index.search(np.array([embedding]), k=top_k)
+        distance, index = distance.flatten().tolist(), index.flatten().tolist()
+        index_url = [WitIndex.wit_dataset['desc2image_map'][i] for i in index]
+        image_info = [WitIndex.wit_dataset['image_info'][i] for i in index_url]
+        print(f"> URL: {image_info[0]}")
+        return distance, index, image_info