added custom download button

README.md

@@ -1,6 +1,6 @@
 ---
 title: VQGAE
-emoji: 📚
+emoji: 🧪
 colorFrom: blue
 colorTo: gray
 sdk: streamlit

app.py

@@ -1,4 +1,9 @@
+import base64
+import os
+import json
 import pickle
+import uuid
+import re
 
 import numpy as np
 import pandas as pd
@@ -14,6 +19,89 @@ from streamlit.components.v1 import html
 MoleculeContainer.depict_settings(aam=False)
 
 
+def download_button(object_to_download, download_filename, button_text, pickle_it=False):
+    """
+    Generates a link to download the given object_to_download.
+    Params:
+    ------
+    object_to_download:  The object to be downloaded.
+    download_filename (str): filename and extension of file. e.g. mydata.csv,
+    some_txt_output.txt download_link_text (str): Text to display for download
+    link.
+    button_text (str): Text to display on download button (e.g. 'click here to download file')
+    pickle_it (bool): If True, pickle file.
+    Returns:
+    -------
+    (str): the anchor tag to download object_to_download
+    Examples:
+    --------
+    download_link(your_df, 'YOUR_DF.csv', 'Click to download data!')
+    download_link(your_str, 'YOUR_STRING.txt', 'Click to download text!')
+    """
+    if pickle_it:
+        try:
+            object_to_download = pickle.dumps(object_to_download)
+        except pickle.PicklingError as e:
+            st.write(e)
+            return None
+
+    else:
+        if isinstance(object_to_download, bytes):
+            pass
+
+        elif isinstance(object_to_download, pd.DataFrame):
+            object_to_download = object_to_download.to_csv(index=False)
+
+        # Try JSON encode for everything else
+        else:
+            object_to_download = json.dumps(object_to_download)
+
+    try:
+        # some strings <-> bytes conversions necessary here
+        b64 = base64.b64encode(object_to_download.encode()).decode()
+
+    except AttributeError as e:
+        b64 = base64.b64encode(object_to_download).decode()
+
+    button_uuid = str(uuid.uuid4()).replace('-', '')
+    button_id = re.sub('\d+', '', button_uuid)
+
+    custom_css = f""" 
+        <style>
+            #{button_id} {{
+                background-color: rgb(255, 255, 255);
+                color: rgb(38, 39, 48);
+                padding: 0.25em 0.38em;
+                position: relative;
+                text-decoration: none;
+                border-radius: 4px;
+                border-width: 1px;
+                border-style: solid;
+                border-color: rgb(230, 234, 241);
+                border-image: initial;
+            }} 
+            #{button_id}:hover {{
+                border-color: rgb(246, 51, 102);
+                color: rgb(246, 51, 102);
+            }}
+            #{button_id}:active {{
+                box-shadow: none;
+                background-color: rgb(246, 51, 102);
+                color: white;
+                }}
+        </style> """
+
+    dl_link = custom_css + f'<a download="{download_filename}" id="{button_id}" href="data:file/txt;base64,{b64}">{button_text}</a><br></br>'
+
+    return dl_link
+
+
+def file_selector(folder_path='.'):
+    filenames = os.listdir(folder_path)
+    selected_filename = st.selectbox('Select a file', filenames)
+    return os.path.join(folder_path, selected_filename)
+
+
 def render_svg(svg_string):
     """Renders the given svg string."""
     c = st.container()
@@ -330,36 +418,45 @@ if submit:
 
     filtered_gen_stats = valid_gen_stats.iloc[filtered_indices]
 
-    with st.form("results"):
-        st.subheader('Generation results', divider='rainbow')
-        st.dataframe(filtered_gen_stats)
-        st.download_button(
-            label="Download results as CSV",
-            data=convert_df(filtered_gen_stats),
-            file_name='vqgae_tubulin_inhibitors_valid.csv',
-            mime='text/csv',
+    st.subheader('Generation results', divider='rainbow')
+    st.dataframe(filtered_gen_stats)
+    download_button(
+        object_to_download=convert_df(filtered_gen_stats),
+        download_filename='vqgae_tubulin_inhibitors_valid.csv',
+        button_text="Download results as CSV"
+    )
+    # st.download_button(
+    #     label="Download results as CSV",
+    #     data=convert_df(filtered_gen_stats),
+    #     file_name='vqgae_tubulin_inhibitors_valid.csv',
+    #     mime='text/csv',
+    # )
+
+    st.subheader('Examples of generated molecules')
+    examples_smiles = filtered_gen_stats.sort_values(by=["similarity_score"], ascending=False).iloc[:6].smiles.to_list()
+    examples = []
+    for smi in examples_smiles:
+        mol = smiles(smi)
+        mol.clean2d()
+        examples.append(mol)
+    svg = grid_depict(examples, 2)
+    render_svg(svg)
+
+    with st.expander("Show full stats"):
+        st.dataframe(full_stats)
+
+        download_button(
+            object_to_download=convert_df(full_stats),
+            download_filename='vqgae_tubulin_inhibitors_full.csv',
+            button_text="Download full results as CSV"
         )
 
-        st.subheader('Examples of generated molecules')
-        examples_smiles = filtered_gen_stats.sort_values(by=["similarity_score"], ascending=False).iloc[:6].smiles.to_list()
-        examples = []
-        for smi in examples_smiles:
-            mol = smiles(smi)
-            mol.clean2d()
-            examples.append(mol)
-        svg = grid_depict(examples, 2)
-        render_svg(svg)
-
-        with st.expander("Show full stats"):
-            st.dataframe(full_stats)
-
-            st.download_button(
-                label="Download full results as CSV",
-                data=convert_df(full_stats),
-                file_name='vqgae_tubulin_inhibitors_full.csv',
-                mime='text/csv',
-            )
-        restart = st.form_submit_button('Restart')
-
-    if restart:
+        # st.download_button(
+        #     label="Download full results as CSV",
+        #     data=convert_df(full_stats),
+        #     file_name='vqgae_tubulin_inhibitors_full.csv',
+        #     mime='text/csv',
+        # )
+
+    if st.button("Restart"):
         st.rerun()