gif-search / app.py
jamescalam's picture
first version
90ae924
raw
history blame contribute delete
No virus
1.7 kB
import streamlit as st
import pinecone
from sentence_transformers import SentenceTransformer
PINECONE_KEY = st.secrets["PINECONE_KEY"] # app.pinecone.io
@st.experimental_singleton(show_spinner=False)
def init_pinecone():
pinecone.init(api_key=PINECONE_KEY, environment="us-west1-gcp")
return pinecone.Index('gif-search')
@st.experimental_singleton(show_spinner=False)
def init_retriever():
return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
with st.spinner("Retrieving the Retriever..."):
retriever = init_retriever()
with st.spinner("Connecting to Pinecone..."):
index = init_pinecone()
def card(urls):
figures = [f"""
<figure style="margin-top: 5px; margin-bottom: 5px; !important;">
<img src="{url}" style="width: 130px; height: 100px; padding-left: 5px; padding-right: 5px" >
</figure>
""" for url in urls]
return st.markdown(f"""
<div style="display: flex; flex-flow: row wrap; text-align: center; justify-content: center;">
{''.join(figures)}
</div>
""", unsafe_allow_html=True)
st.write("""
## ⚡️ AI-Powered GIF Search ⚡️
Search for GIFs using Semantic Search, learn how it works [here](https://www.pinecone.io/learn/gif-search/).
""")
query = st.text_input("What are you looking for?", "")
if query != "":
with st.spinner(text="Similarity Searching..."):
xq = retriever.encode([query]).tolist()
xc = index.query(xq, top_k=30, include_metadata=True)
urls = []
for context in xc['matches']:
urls.append(context['metadata']['url'])
with st.spinner(text="Fetching GIFs 🚀🚀🚀"):
card(urls)