Fix 69 - Refactor how arguments are added to scripts

examples/blend_and_shuffle.py

@@ -17,7 +17,7 @@
 import nemo_curator as nc
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def main(args):
@@ -28,7 +28,7 @@ def main(args):
     output_path = "/path/to/output"
 
     # Set up Dask client
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     # Blend the datasets
     datasets = [DocumentDataset.read_json(path) for path in dataset_paths]
@@ -46,7 +46,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/classifier_filtering.py

@@ -23,7 +23,7 @@
 from nemo_curator.modifiers import FastTextLabelModifier
 from nemo_curator.utils.distributed_utils import get_client, read_data, write_to_disk
 from nemo_curator.utils.file_utils import get_all_files_paths_under
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def load_dataset(input_data_dir):
@@ -55,7 +55,7 @@ def main(args):
     filtered_output = "/path/to/output"
 
     # Prepare samples for the classifier
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
     low_quality_samples = create_samples(
         low_quality_data_path, "__label__lq", num_low_quality_samples
     )
@@ -100,7 +100,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/domain_classifier_example.py

@@ -19,7 +19,7 @@
 from nemo_curator import DomainClassifier
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def main(args):
@@ -60,7 +60,7 @@ def main(args):
     input_file_path = "/path/to/data"
     output_file_path = "./"
 
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     input_dataset = DocumentDataset.read_json(
         input_file_path, backend="cudf", add_filename=True
@@ -89,54 +89,18 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    parser.add_argument(
-        "--scheduler-address",
-        type=str,
-        default=None,
-        help="Address to the scheduler of a created dask cluster. If not provided"
-        "a single node LocalCUDACluster will be started.",
-    )
-    parser.add_argument(
-        "--scheduler-file",
-        type=str,
-        default=None,
-        help="Path to the scheduler file of a created dask cluster. If not provided"
-        " a single node LocalCUDACluster will be started.",
-    )
-    parser.add_argument(
-        "--nvlink-only",
-        action="store_true",
-        help="Start a local cluster with only NVLink enabled."
-        "Only applicable when protocol=ucx and no scheduler file/address is specified",
-    )
-    parser.add_argument(
-        "--protocol",
-        type=str,
-        default="ucx",
-        help="Protcol to use for dask cluster"
-        "Note: This only applies to the localCUDACluster. If providing an user created "
-        "cluster refer to"
-        "https://docs.rapids.ai/api/dask-cuda/stable/api.html#cmdoption-dask-cuda-protocol",  # noqa: E501
-    )
-    parser.add_argument(
-        "--rmm-pool-size",
-        type=str,
-        default="14GB",
-        help="Initial pool size to use for the RMM Pool Memory allocator"
-        "Note: This only applies to the localCUDACluster. If providing an user created "
-        "cluster refer to"
-        "https://docs.rapids.ai/api/dask-cuda/stable/api.html#cmdoption-dask-cuda-rmm-pool-size",  # noqa: E501
-    )
-    parser.add_argument("--enable-spilling", action="store_true")
-    parser.add_argument("--set-torch-to-use-rmm", action="store_true")
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="gpu",
-        help="Device to run the script on. Either 'cpu' or 'gpu'.",
-    )
-
-    return parser
+    argumentHelper = ArgumentHelper(parser)
+
+    argumentHelper.add_arg_device()
+    argumentHelper.add_arg_enable_spilling()
+    argumentHelper.add_arg_nvlink_only()
+    argumentHelper.add_arg_protocol()
+    argumentHelper.add_arg_rmm_pool_size()
+    argumentHelper.add_arg_scheduler_address()
+    argumentHelper.add_arg_scheduler_file()
+    argumentHelper.add_arg_set_torch_to_use_rmm()
+
+    return argumentHelper.parser
 
 
 if __name__ == "__main__":

examples/download_arxiv.py

@@ -16,7 +16,7 @@
 
 from nemo_curator.download import download_arxiv
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def main(args):
@@ -27,7 +27,7 @@ def main(args):
     url_limit = 10
 
     # Set up Dask client
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     # Download and sample data
     arxiv = download_arxiv(output_directory, url_limit=url_limit)
@@ -42,7 +42,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/download_common_crawl.py

@@ -16,7 +16,7 @@
 
 from nemo_curator.download import download_common_crawl
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def main(args):
@@ -29,7 +29,7 @@ def main(args):
     url_limit = 10
 
     # Set up Dask client
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     # Download and sample data
     common_crawl = download_common_crawl(
@@ -46,7 +46,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/download_wikipedia.py

@@ -16,7 +16,7 @@
 
 from nemo_curator.download import download_wikipedia
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def main(args):
@@ -28,7 +28,7 @@ def main(args):
     url_limit = 10
 
     # Set up Dask client
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     # Download and sample data
     wikipedia = download_wikipedia(
@@ -45,7 +45,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/exact_deduplication.py

@@ -19,7 +19,7 @@
 from nemo_curator.modules import ExactDuplicates
 from nemo_curator.utils.distributed_utils import get_client, read_data, write_to_disk
 from nemo_curator.utils.file_utils import get_all_files_paths_under
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def pre_imports():
@@ -33,7 +33,7 @@ def main(args):
     output_dir = "./"
     dataset_id_field = "id"
     dataset_text_field = "text"
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
     backend = "cudf" if args.device == "gpu" else "pandas"
 
     if args.device == "gpu":
@@ -79,7 +79,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/find_pii_and_deidentify.py

@@ -21,13 +21,13 @@
 from nemo_curator.modifiers.pii_modifier import PiiModifier
 from nemo_curator.modules.modify import Modify
 from nemo_curator.utils.distributed_utils import get_client
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def console_script():
     parser = argparse.ArgumentParser()
-    arguments = add_distributed_args(parser).parse_args()
-    _ = get_client(**parse_client_args(arguments))
+    args = ArgumentHelper(parser).add_distributed_args().parse_args()
+    _ = get_client(**ArgumentHelper.parse_client_args(args))
 
     dataframe = pd.DataFrame(
         {"text": ["Sarah and Ryan went out to play", "Jensen is the CEO of NVIDIA"]}

examples/fuzzy_deduplication.py

@@ -21,7 +21,7 @@
 from nemo_curator import FuzzyDuplicates, FuzzyDuplicatesConfig
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.utils.distributed_utils import get_client, write_to_disk
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def pre_imports():
@@ -44,7 +44,7 @@ def main(args):
     assert args.device == "gpu"
 
     with dask.config.set({"dataframe.backend": backend}):
-        client = get_client(**parse_client_args(args))
+        client = get_client(**ArgumentHelper.parse_client_args(args))
         client.run(pre_imports)
 
         t0 = time.time()
@@ -102,7 +102,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/identify_languages_and_fix_unicode.py

@@ -24,7 +24,7 @@
     get_all_files_paths_under,
     separate_by_metadata,
 )
-from nemo_curator.utils.script_utils import add_distributed_args, parse_client_args
+from nemo_curator.utils.script_utils import ArgumentHelper
 
 
 def load_dataset(input_data_dir):
@@ -49,7 +49,7 @@ def main(args):
     language_field = "language"
 
     # Prepare samples for the classifier
-    client = get_client(**parse_client_args(args))
+    client = get_client(**ArgumentHelper.parse_client_args(args))
 
     # Filter data
     multilingual_dataset = load_dataset(multilingual_data_path)
@@ -88,7 +88,7 @@ def attach_args(
         formatter_class=argparse.ArgumentDefaultsHelpFormatter
     ),
 ):
-    return add_distributed_args(parser)
+    return ArgumentHelper(parser).add_distributed_args()
 
 
 if __name__ == "__main__":

examples/k8s/create_dask_cluster.py

@@ -2,6 +2,8 @@
 
 from dask_kubernetes.operator.kubecluster import KubeCluster, make_cluster_spec
 
+from nemo_curator.utils.script_utils import ArgumentHelper
+
 
 def create_cluster(
     name: str,
@@ -127,8 +129,6 @@ def parse_pvcs(specs: str) -> dict[str, str]:
         help="Comma sep PVC specificiation of $pvc_name_1:$mount_path_1,$pvc_name_2:$mount_path_2. Example: foo:/foo,bar:/bar mounts pvcs named foo and bar to /foo and /bar respectively.",
     )
 
-    args = parser.parse_args()
-
     create_cluster(
-        **vars(args),
+        **vars(parser.parse_args()),
     )