Update app.py

app.py

@@ -7,7 +7,11 @@ import torch
 import transformers
 from transformers import RobertaModel, AutoTokenizer
 
-@st.cache(hash_funcs={transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast: hash}, suppress_st_warning=True, allow_output_mutation=True)
+#@st.cache(hash_funcs={transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast: hash}, suppress_st_warning=True, allow_output_mutation=True)
+@st.cache(show_spinner=False,
+          hash_funcs={RobertaModel: lambda _: None,
+                      AutoTokenizer: lambda _: None,
+                      dict: lambda _: None})
 def load():
     text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text')
     tokenizer = AutoTokenizer.from_pretrained('SajjadAyoubi/clip-fa-text')
@@ -28,12 +32,12 @@ def get_html(url_list, height=224):
     return HTML
 
 def compute_embeddings(query):
-    return text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
+    with torch.no_grad():
+        return text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
 
-@st.cache(show_spinner=False)
+st.cache(show_spinner=False)
 def image_search(query, top_k=8):
-    with torch.no_grad():
-        text_embedding = compute_embeddings(query)
+    text_embedding = compute_embeddings(query)
     values, indices = torch.cosine_similarity(text_embedding, image_embeddings).sort(descending=True)
     return [links[i] for i in indices[:top_k]]