WIP

script.py

@@ -52,8 +52,11 @@ if __name__ == "__main__":
 
     if data_path_test_server.exists():
         data_path = data_path_test_server
+        TEST_ENV = True
     else: 
         data_path = data_path_local
+        TEST_ENV = False
+
     
     print(data_path)
 
@@ -65,22 +68,23 @@ if __name__ == "__main__":
     print(data_files)
 
 
-    try:
+    if TEST_ENV:
         dataset = load_dataset(
-            "arrow", 
+            "webdataset", 
             data_files=data_files,
             trust_remote_code=True,
             # streaming=True
         )
-        print('load with arrow')
-    except:
+        print('load with webdataset')
+    else:
         dataset = load_dataset(
-            "webdataset", 
+            "arrow", 
             data_files=data_files,
             trust_remote_code=True,
             # streaming=True
         )
-        print('load with webdataset')
+        print('load with arrow')
+    
 
     print(dataset, flush=True)
     # dataset = load_dataset('webdataset', data_files={)