WIP

script.py

@@ -56,18 +56,30 @@ if __name__ == "__main__":
         data_path = data_path_local
     
     print(data_path)
-    print([str(p) for p in data_path.rglob('*validation*.arrow')])
+    print([str(p) for p in data_path.rglob('*validation*.(arrow|tar)')])
 
     # dataset = load_dataset(params['dataset'], trust_remote_code=True, use_auth_token=params['token'])
-    dataset = load_dataset(
-        "arrow", 
-        data_files={
-            "validation": [str(p) for p in data_path.rglob('*validation*.arrow')],
-            "test": [str(p) for p in data_path.rglob('*test*.arrow')],
-        },
-        trust_remote_code=True,
-        # streaming=True
-    )
+    data_files = {
+        "validation": [str(p) for p in [*data_path.rglob('*validation*.arrow')]+[*data_path.rglob('*validation*.tar')]],
+        "test": [str(p) for p in [*data_path.rglob('*test*.arrow')]+[*data_path.rglob('*test*.tar')]],
+    }
+    try:
+        dataset = load_dataset(
+            "arrow", 
+            data_files=data_files,
+            trust_remote_code=True,
+            # streaming=True
+        )
+        print('load with arrow')
+    except:
+        dataset = load_dataset(
+            "webdataset", 
+            data_files=data_files,
+            trust_remote_code=True,
+            # streaming=True
+        )
+        print('load with webdataset')
+
     print(dataset, flush=True)
     # dataset = load_dataset('webdataset', data_files={)