Update streamlitapp.py

streamlitapp.py

@@ -8,15 +8,15 @@ import streamlit.components.v1 as components
 from datasets import load_dataset
 
 proj_dir = Path(__file__).parent
-
 servername = 'comet-team-kangas-demo.hf.space'
 src = f"https://{servername}/kangas/?datagrid=.%2Fdatagrids%2Fcoco-500.datagrid"
-
 st.set_page_config(layout="wide")
 
-st.markdown("1. Select dataset of your choice")
+st.markdown("# HuggingFace + Streamlit + Kangas")
+st.markdown("1. Select a HuggingFace dataset, or select one below")
 
 def kangas_fn(dataset_repo):
+    global src
     dg_file_name = dataset_repo.replace('/', '__') + '.datagrid'
     dg_full_path = str(proj_dir / 'datagrids' / dg_file_name)
     with st.spinner("2. Loading Dataset..."):
@@ -24,11 +24,13 @@ def kangas_fn(dataset_repo):
     with st.spinner("3. Creating Kangas..."):
         if not os.path.exists(dg_full_path):
             dg = kg.DataGrid(dataset)
+    with st.spinner("4. Saving Kangas..."):
+        if not os.path.exists(dg_full_path):
             dg.save(dg_full_path)
     with st.spinner("4. Loading Kangas..."):
         src = f"https://{servername}/kangas/?datagrid=.%2Fdatagrids%2F" + dg_file_name
         
 hf_dataset = st.text_input("HuggingFace Dataset", value='beans')
-st.button("Download and Run", on_click=partial(kangas_fn, hf_dataset))
-st.markdown("""Click the dropdown in Kangas to see pre-loaded datasets""")
+st.button("Download", on_click=partial(kangas_fn, hf_dataset))
+st.markdown("""Click the dropdown in Kangas select a dataset""")
 components.iframe(src, None, 1000, scrolling=True)