Update src/streamlit_app.py

src/streamlit_app.py

@@ -52,10 +52,16 @@ clip_model, clip_processor, text_model = load_models()
 
 # ========== 📥 Load Dataset ==========
 @st.cache_resource(show_spinner=False)
+from datasets import get_dataset_split_names
+
+available_splits = get_dataset_split_names("univanxx/3mdbench")
+# Pick 'train' if available, otherwise first available split
+split_to_use = "train" if "train" in available_splits else available_splits[0]
+
 def load_medical_data():
     dataset = load_dataset(
         "univanxx/3mdbench",
-        split="train",
+        split=split_to_use,
         cache_dir=os.environ["HF_DATASETS_CACHE"]
     )
     return dataset