Update app.py

app.py

@@ -9,7 +9,6 @@ import numpy as np
 import os
 import urllib.request
 import zipfile
-import collections
 
 
 class MLP(torch.nn.Module):
@@ -146,7 +145,7 @@ def find_datasets_and_model_ids(root_dir):
                     dataset_name = os.path.basename(os.path.dirname(stoke_config_path))
                     datasets.setdefault(model_id, {})[dataset_name] = stoke_config_data
     
-    return collections.OrderedDict(sorted(datasets.items()))
+    return datasets
 
 
 # Main content
@@ -166,7 +165,7 @@ with st.sidebar:
     # Select dataset based on selected model
     available_datasets = datasets[model_selection]
     if available_datasets:
-        dataset_selection = st.selectbox("Select Dataset", available_datasets)
+        dataset_selection = st.selectbox("Select Dataset", sorted(available_datasets))
     else:
         st.error("No datasets available for the selected model.")