Skip to content

Commit a902456

Browse files
Add num_proc to CLI test
1 parent 4693108 commit a902456

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

src/datasets/commands/test.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ def _test_command_factory(args):
2828
args.ignore_verifications,
2929
args.force_redownload,
3030
args.clear_cache,
31+
args.num_proc,
3132
)
3233

3334

@@ -65,6 +66,7 @@ def register_subcommand(parser: ArgumentParser):
6566
action="store_true",
6667
help="Remove downloaded files and cached datasets after each config test",
6768
)
69+
test_parser.add_argument("--num_proc", type=int, default=None, help="Number of processes")
6870
# aliases
6971
test_parser.add_argument("--save_infos", action="store_true", help="alias to save_info")
7072
test_parser.add_argument("dataset", type=str, help="Name of the dataset to download")
@@ -81,6 +83,7 @@ def __init__(
8183
ignore_verifications: bool,
8284
force_redownload: bool,
8385
clear_cache: bool,
86+
num_proc: int,
8487
):
8588
self._dataset = dataset
8689
self._name = name
@@ -91,6 +94,7 @@ def __init__(
9194
self._ignore_verifications = ignore_verifications
9295
self._force_redownload = force_redownload
9396
self._clear_cache = clear_cache
97+
self._num_proc = num_proc
9498
if clear_cache and not cache_dir:
9599
print(
96100
"When --clear_cache is used, specifying a cache directory is mandatory.\n"
@@ -151,6 +155,7 @@ def get_builders() -> Generator[DatasetBuilder, None, None]:
151155
if self._ignore_verifications
152156
else VerificationMode.ALL_CHECKS,
153157
try_from_hf_gcs=False,
158+
num_proc=self._num_proc,
154159
)
155160
builder.as_dataset()
156161
if self._save_infos:

0 commit comments

Comments
 (0)