Update app.py

app.py

@@ -1,24 +1,21 @@
 import streamlit as st
-import pandas as pd, numpy as np
+import pandas as pd
 from html import escape
 import os
 import torch
 from transformers import RobertaModel, AutoTokenizer
 
 
-@st.cache(show_spinner=False,
-          hash_funcs={text_encoder: lambda _: None,
-                      tokenizer: lambda _: None,
-                      dict: lambda _: None})
+@st.cache(show_spinner=False)
 def load():
     text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text')
     tokenizer = AutoTokenizer.from_pretrained('SajjadAyoubi/clip-fa-text')
-    df = pd.read_csv('data.csv')
-    image_embeddings = np.load('embeddings.npy')
-    return text_encoder, tokenizer, df, image_embeddings
+    link_df = pd.read_csv('links.csv')
+    image_embeddings = torch.load('embeddings.pt')
+    return text_encoder, tokenizer, link_df, image_embeddings
 
 
-text_encoder, tokenizer, df, image_embeddings = load()
+text_encoder, tokenizer, link_df, image_embeddings = load()
 
 
 def get_html(url_list, height=224):
@@ -34,12 +31,12 @@ def get_html(url_list, height=224):
     return html
 
 
-st.cache(show_spinner=False)
+@st.cache(show_spinner=False)
 def image_search(query, top_k=8):
     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 [(df.iloc[i]['path'], df.iloc[i]['link']) for i in indices[:top_k]]
+    return [(link_df.iloc[i]['path'], link_df.iloc[i]['link']) for i in indices[:top_k]]
 
 
 description = '''