Initial code

.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db filter=lfs diff=lfs merge=lfs -text
+*.faiss filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

.github/workflows/hf_sync.yml

@@ -0,0 +1,20 @@
+name: Sync to Hugging Face hub
+on:
+  push:
+    branches: [main]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push --force https://Tuana:[email protected]/spaces/Tuana/find-the-animal main

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+.DS_Store

README.md

@@ -0,0 +1,10 @@
+---
+title: MultiModalRetrival for Image Search
+emoji: 😽
+colorFrom: green
+colorTo: yellow
+sdk: streamlit
+sdk_version: 1.2.0
+app_file: app.py
+pinned: false
+---

app.py

@@ -0,0 +1,80 @@
+import streamlit as st
+import time
+import logging
+from json import JSONDecodeError
+from PIL import Image
+
+
+from utils.haystack import query
+from utils.frontend import reset_results, set_state_if_absent
+
+def main():
+
+    set_state_if_absent("statement", "What is the fastest animal?")
+    set_state_if_absent("results", None)
+
+    st.write("# Look for images with MultiModalRetrieval 🐅")
+    st.write()
+    st.markdown(
+        """
+    ##### Enter a question about animals
+    """
+    )
+    # Search bar
+    statement = st.text_input(
+        "", value=st.session_state.statement, max_chars=100, on_change=reset_results
+    )
+        
+    col1, col2 = st.columns(2)
+    col1.markdown(
+        "<style>.stButton button {width:100%;}</style>", unsafe_allow_html=True
+    )
+   
+    run_pressed = col1.button("Run")
+
+    run_query = (
+        run_pressed or statement != st.session_state.statement
+    )
+
+    # Get results for query
+    if run_query and statement:
+        time_start = time.time()
+        reset_results()
+        st.session_state.statement = statement
+        with st.spinner("🧠 &nbsp;&nbsp; Performing neural search on documents..."):
+            try:
+                docs = query(statement)
+                st.write(docs["documents"])
+                for doc in docs["documents"]:
+                    image = Image.open(doc.content)
+                    st.image(image)
+                for answer in docs["answers"]:
+                    st.write(answer)
+                print(f"S: {statement}")
+                time_end = time.time()
+                print(time.strftime("%Y-%m-%d %H:%M:%S", time.gmtime()))
+                print(f"elapsed time: {time_end - time_start}")
+            except JSONDecodeError as je:
+                st.error(
+                    "👓 &nbsp;&nbsp; An error occurred reading the results. Is the document store working?"
+                )
+                return
+            except Exception as e:
+                logging.exception(e)
+                st.error("🐞 &nbsp;&nbsp; An error occurred during the request.")
+            return
+    
+    # if st.session_state.results:
+    #     st.write("Got some results")
+    #     print("GOT RESTULTS")
+        # st.write("Received Results")
+        # results = st.session_state.results
+        # print(results)
+        # docs = results["documents"]
+        # st.write(results)
+        # # show different messages depending on entailment results
+        # for doc in docs:
+        #     image = Image(filename=doc.content)
+        #     st.image(image)
+
+main()

images.db

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1df1fbb3cd45b562b6561acc3b159ea637ee538bc2f8ce2c59fa959dbc7b2538
+size 200704

requirements.txt

@@ -0,0 +1,2 @@
+farm-haystack[faiss]==1.11.1
+streamlit==1.12.0

text.db

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fa05449dc61ca9bb83ad4053fb8797eb0b4deba7a31e2b2b15d0d73f97c3095
+size 4464640

utils/config.py

@@ -0,0 +1 @@
+INDEX_DIR = "data/index"

utils/frontend.py

@@ -0,0 +1,10 @@
+import streamlit as st
+
+def set_state_if_absent(key, value):
+    if key not in st.session_state:
+        st.session_state[key] = value
+
+def reset_results(*args):
+    st.write("Called reset")
+    st.session_state.answer = None
+    st.session_state.results = None

utils/haystack.py

@@ -0,0 +1,79 @@
+import shutil
+from haystack.document_stores import FAISSDocumentStore
+from haystack.nodes.retriever import EmbeddingRetriever, MultiModalRetriever
+from haystack.nodes.reader import FARMReader
+from haystack import Pipeline
+from utils.config import (INDEX_DIR)
+from typing import List
+from haystack import BaseComponent, Answer
+import streamlit as st
+
+
+
+class AnswerToQuery(BaseComponent):
+
+  outgoing_edges = 1
+
+  def run(self, query: str, answers: List[Answer]):
+    return {"query": answers[0].answer}, "output_1"
+
+  def run_batch(self):
+    raise NotImplementedError()
+
+# cached to make index and models load only at start
+@st.cache(
+    hash_funcs={"builtins.SwigPyObject": lambda _: None}, allow_output_mutation=True
+)
+def start_haystack():
+    """
+    load document store, retriever, entailment checker and create pipeline
+    """
+    shutil.copy(f"{INDEX_DIR}/text.db", ".")
+    shutil.copy(f"{INDEX_DIR}/images.db", ".")
+
+    document_store_text = FAISSDocumentStore(
+        faiss_index_path=f"{INDEX_DIR}/text.faiss",
+        faiss_config_path=f"{INDEX_DIR}/text.json",
+    )
+    
+    document_store_images = FAISSDocumentStore(
+        faiss_index_path=f"{INDEX_DIR}/images.faiss",
+        faiss_config_path=f"{INDEX_DIR}/images.json",
+    )
+    retriever_text = EmbeddingRetriever(
+        document_store=document_store_text,
+        embedding_model="sentence-transformers/multi-qa-mpnet-base-dot-v1",
+        model_format="sentence_transformers",
+    )
+
+    reader = FARMReader(model_name_or_path="deepset/roberta-base-squad2", use_gpu=True)
+    
+
+    retriever_images = MultiModalRetriever(
+        document_store=document_store_images,
+        query_embedding_model = "sentence-transformers/clip-ViT-B-32",
+        query_type="text",
+        document_embedding_models = {
+            "image": "sentence-transformers/clip-ViT-B-32"
+        }
+    )
+    
+    answer_to_query = AnswerToQuery()
+
+    pipe = Pipeline()
+
+    pipe.add_node(retriever_text, name="text_retriever", inputs=["Query"])
+    pipe.add_node(reader, name="text_reader", inputs=["text_retriever"])
+    pipe.add_node(answer_to_query, name="answer2query", inputs=["text_reader"])
+    pipe.add_node(retriever_images, name="image_retriever", inputs=["answer2query"])
+
+    return pipe
+    
+pipe = start_haystack()
+
+@st.cache(allow_output_mutation=True)
+def query(statement: str, text_retriever_top_k: int = 5, image_retriever_top_k = 1):
+    """Run query and verify statement"""
+    params = {"image_retriever": {"top_k": image_retriever_top_k},"text_retriever": {"top_k": text_retriever_top_k} }
+    results = pipe.run(statement, params=params)
+    return results