Update app.py

app.py

@@ -4,9 +4,14 @@ from datasets import load_dataset
 
 auth_token = os.environ.get("auth_token")
 visit_bench_all = load_dataset("mlfoundations/visit-bench", use_auth_token=auth_token)
+print('visit_bench_all')
+print(visit_bench_all)
 print('dataset keys:')
 print(visit_bench_all.keys())
-visit_bench = visit_bench_all['test']
+dataset_keys = list(visit_bench_all.keys())
+assert len(dataset_keys) == 1
+dataset_key = dataset_keys[0]
+visit_bench = visit_bench_all[dataset_key]
 
 df = visit_bench.to_pandas()
 print(f"Got {len(df)} items in dataframe")