fix

app.py

@@ -17,7 +17,7 @@ def average_pool(last_hidden_states: Tensor,
     return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
         
         
-@st.cache(allow_output_mutation=True)
+@st.cache_resource
 def load_model_and_tokenizer():
     tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large')
     model = AutoModel.from_pretrained('intfloat/multilingual-e5-large')
@@ -26,14 +26,14 @@ def load_model_and_tokenizer():
     return model, tokenizer
 
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource
 def load_title_data():
     title_df = pd.read_csv('anlp2024.tsv', names=["pid", "title"], sep="\t")
     
     return title_df
 
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource
 def load_title_embeddings():
     npz_comp = np.load("anlp2024.npz")
     title_embeddings = npz_comp["arr_0"]
@@ -41,13 +41,13 @@ def load_title_embeddings():
     return title_embeddings
 
 
-@st.cache
+@st.cache_data
 def get_retrieval_results(index, input_text, top_k, tokenizer, title_df):
     batch_dict = tokenizer(f"query: {input_text}", max_length=512, padding=True, truncation=True, return_tensors='pt')
     with torch.no_grad():
         outputs = model(**batch_dict)
-        embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
-        embeddings = F.normalize(embeddings, p=2, dim=1)
+        query_embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+        query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
     
     _, ids = index.search(x=np.array([query_embeddings]), k=top_k)
     retrieved_titles = []