Update app.py

app.py

@@ -32,16 +32,9 @@ st.cache(show_spinner=False)
 def image_search(query, top_k=10):
     with torch.no_grad():
         text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
-    values, indices = torch.cosine_similarity(text_embedding, image_embeddings).sort(descending=True)
+    values, indices = torch.dot(image_embeddings, text_embedding.T).sort(descending=True)
     return [links[i] for i in indices[:top_k]]
 
-st.cache(show_spinner=False)     
-def image_search_(query, top_k=10):
-    with torch.no_grad():
-        text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output.numpy()
-    results = np.argsort((image_embeddings@text_embedding.T)[:, 0])[-1:-top_k-1:-1]
-    return [links[i] for i in results]
-
 
 description = '''
 # Persian (fa) image search 
@@ -79,7 +72,7 @@ def main():
     _, c, _ = st.columns((1, 3, 1))
     query = c.text_input('Search Box (type in fa)', value='گل صورتی')
     if len(query) > 0:
-        results = image_search_(query)
+        results = image_search(query)
         st.markdown(get_html(results), unsafe_allow_html=True)