Update app.py

app.py

@@ -17,6 +17,13 @@ tokenizer, text_encoder, links, image_embeddings = load()
 
 
 @st.experimental_memo
+def image_search(query, top_k=10):
+    with torch.no_grad():
+        text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
+    _, indices = torch.cosine_similarity(image_embeddings, text_embedding).sort(descending=True)
+    return [links[i] for i in indices[:top_k]]
+
+
 def get_html(url_list):
     html = "<div style='margin-top: 50px; max-width: 1100px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>"
     for url in url_list:
@@ -24,13 +31,6 @@ def get_html(url_list):
         html = html + html2
     html += "</div>"
     return html
- 
-@st.experimental_memo
-def image_search(query, top_k=10):
-    with torch.no_grad():
-        text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
-    _, indices = torch.cosine_similarity(image_embeddings, text_embedding).sort(descending=True)
-    return [links[i] for i in indices[:top_k]]
 
 
 description = '''
@@ -68,6 +68,7 @@ def main():
     st.sidebar.markdown(description)
     _, c, _ = st.columns((1, 3, 1))
     query = c.text_input('Search Box (type in fa)', value='گل صورتی')
+    c.text("It'll take about 20s to load all new images")
     if len(query) > 0:
         results = image_search(query)
         st.markdown(get_html(results), unsafe_allow_html=True)