bugfix: remove workers from load_from_disk

SeanLee97 · Jan 11, 2025 · d007d52 · d007d52
1 parent 18af7f3
commit d007d52
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/angle_emb/angle_trainer.py b/angle_emb/angle_trainer.py
@@ -211,7 +211,7 @@ def main():
 
     if os.path.exists(args.train_name_or_path):
         if os.path.isdir(args.train_name_or_path):
-            ds = load_from_disk(args.train_name_or_path, num_proc=args.workers)
+            ds = load_from_disk(args.train_name_or_path)
         else:
             ds = load_dataset('json',
                               data_files=[args.train_name_or_path],
@@ -244,7 +244,7 @@ def main():
         logger.info('Validation detected, processing validation...')
         if os.path.exists(args.valid_name_or_path):
             if os.path.isdir(args.valid_name_or_path):
-                valid_ds = load_from_disk(args.valid_name_or_path, num_proc=args.workers)
+                valid_ds = load_from_disk(args.valid_name_or_path)
             else:
                 valid_ds = load_dataset('json', data_files=[args.valid_name_or_path], num_proc=args.workers)
         else:
@@ -263,7 +263,7 @@ def main():
         logger.info('Validation for callback detected, processing validation...')
         if os.path.exists(args.valid_name_or_path_for_callback):
             if os.path.isdir(args.valid_name_or_path_for_callback):
-                valid_ds_for_callback = load_from_disk(args.valid_name_or_path_for_callback, num_proc=args.workers)
+                valid_ds_for_callback = load_from_disk(args.valid_name_or_path_for_callback)
             else:
                 valid_ds_for_callback = load_dataset(
                     'json', data_files=[args.valid_name_or_path_for_callback], num_proc=args.workers)