final download option

app.py

@@ -52,12 +52,6 @@ def about_page():
     )
 
     st.image('figures/hyper-dti.png', caption='Overview of HyperPCM architecture.')
-
-
-@st.cache
-def convert_df(df):
-    # IMPORTANT: Cache the conversion to prevent computation on every rerun
-    return df.to_csv().encode('utf-8')
     
 
 def retrieval():
@@ -161,9 +155,7 @@ def retrieval():
                     mol_img = Chem.Draw.MolToImage(mol)
                     st.image(mol_img, caption=f"{results.loc[j + 5*i, 'Prediction']:.2f}")
 
-        st.download_button('Download retrieved drug compounds.', convert_df(results.head(selected_k)), file_name='retrieved_drugs.csv')
-
-        st.markdown('## TMP title')
+        st.download_button('Download retrieved drug compounds.', results.head(selected_k).to_csv().encode('utf-8'), file_name='retrieved_drugs.csv')
                     
     elif query_embedding is not None:
         st.markdown('### Inference')
@@ -213,9 +205,7 @@ def retrieval():
                     mol_img = Chem.Draw.MolToImage(mol)
                     st.image(mol_img, caption=f"{results.loc[j + 5*i, 'Prediction']:.2f}")
 
-        st.download_button('Download retrieved drug compounds.', convert_df(results.head(selected_k)), file_name='retrieved_drugs.csv')
-
-        st.markdown('## TMP title')
+        st.download_button('Download retrieved drug compounds.', results.head(selected_k).to_csv().encode('utf-8'), file_name='retrieved_drugs.csv')
     
 page_names_to_func = {
     'Retrieval': retrieval,