Show a please wait message when caching datasets

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import streamlit_pianoroll
 from fortepyan import MidiPiece
 
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 
 
 def main():
@@ -12,6 +12,7 @@ def main():
         "giant-midi-sustain-v2",
         "maestro-sustain-v2",
         "imslp-midi-v1",
+        "atepp-1.1-sustain-v2",
     ]
     dataset_name = st.selectbox(
         label="Select dataset",
@@ -20,10 +21,17 @@ def main():
     preview_dataset(dataset_name)
 
 
-def preview_dataset(dataset_name: str):
+@st.cache_data(persist="disk")
+def get_dataset(dataset_name: str) -> Dataset:
     dataset = load_dataset(f"epr-labs/{dataset_name}", split="train")
 
+    return dataset
+
+
+def preview_dataset(dataset_name: str):
     st.write(f"### Dataset: {dataset_name}")
+    with st.spinner("Loading the dataset, hang tight!"):
+        dataset = get_dataset(dataset_name)
 
     st.write(f"""
     Number of records: {len(dataset)}