Update app.py

app.py

@@ -5,11 +5,18 @@ import torch
 from transformers import RobertaModel, AutoTokenizer
 
 
-tokenizer = AutoTokenizer.from_pretrained('SajjadAyoubi/clip-fa-text')
-text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text').eval()
-image_embeddings = torch.load('embedding.pt')
-links = np.load('data.npy', allow_pickle=True)
-
+@st.cache(show_spinner=False,
+          hash_funcs={AutoTokenizer: lambda _: None,
+                      RobertaModel: lambda _: None,
+                      dict: lambda _: None})
+def load():
+  tokenizer = AutoTokenizer.from_pretrained('SajjadAyoubi/clip-fa-text')
+  text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text').eval()
+  image_embeddings = torch.load('embedding.pt').numpy()
+  links = np.load('data.npy', allow_pickle=True)
+  return model, processor, df, embeddings
+  
+model, processor, df, embeddings = load()
 
 
 def get_html(url_list):
@@ -20,19 +27,26 @@ def get_html(url_list):
     html += "</div>"
     return html
  
-                           
+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)
     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 
 - Enter your query and hit enter
 
-Built with [CLIP-fa](https://github.com/sajjjadayobi/CLIPfa) model and 25k images from Unsplash
+Built with [CLIP-fa](https://github.com/sajjjadayobi/CLIPfa) model and 25k images from [Unsplash](https://unsplash.com/)
 '''
 
 
@@ -62,9 +76,9 @@ def main():
                 
     st.sidebar.markdown(description)
     _, c, _ = st.columns((1, 3, 1))
-    query = c.text_input('Search Box', value='گل صورتی')
+    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)