update

app.py

@@ -65,7 +65,7 @@ def build_faiss_index(sentence_emeddings):
     return index
 
 
-@st.cache
+@st.cache(allow_output_mutation=True)
 def get_retrieval_results(index, input_text, top_k, model, tokenizer, sentence_df):
     with torch.no_grad():
         inputs = tokenizer.encode_plus(
@@ -80,20 +80,23 @@ def get_retrieval_results(index, input_text, top_k, model, tokenizer, sentence_d
         query_embeddings = query_embeddings.detach().cpu().numpy()
         query_embeddings = query_embeddings / np.linalg.norm(query_embeddings, ord=2)
 
-    print(np.array([query_embeddings]))
+    _, ids = index.search(x=np.array([query_embeddings]), k=top_k)
+    retrieved_sentences = []
+
+    for id in ids[0]:
+        retrieved_sentences.append(sentence_df.loc[id, "sentence"])
 
-    dists, ids = index.search(x=np.array([query_embeddings]), k=top_k)
-    print(dists)
-    print(ids)
+    return pd.DataFrame({"sentences": retrieved_sentences})
 
 
-def main(model, tokenizer, sentence_df, sentence_embeddings, index):
+def main(model, tokenizer, sentence_df, index):
     st.markdown("## AI-based Paraphrasing for Academic Writing")
 
     input_text = st.text_area("text input", "Model have good results.", placeholder="Write something here...")
     top_k = st.number_input('top_k', min_value=1, value=10, step=1)
 
-    get_retrieval_results(index, input_text, top_k, model, tokenizer, sentence_df)
+    df = get_retrieval_results(index, input_text, top_k, model, tokenizer, sentence_df)
+    st.table(df)
 
 
 if __name__ == "__main__":
@@ -104,4 +107,4 @@ if __name__ == "__main__":
     faiss.normalize_L2(sentence_emeddings)
     index = build_faiss_index(sentence_emeddings)
 
-    main(model, tokenizer, sentence_df, sentence_emeddings, index)
+    main(model, tokenizer, sentence_df, index)