Add AEGIS classifier

.pre-commit-config.yaml

@@ -33,6 +33,7 @@ repos:
         exclude: docs/
       - id: requirements-txt-fixer
       - id: trailing-whitespace
+        exclude: nemo_curator/utils/aegis_utils.py
 
   - repo: https://github.com/psf/black
     rev: 24.4.2

examples/classifiers/aegis_example.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import time
+
+from nemo_curator.classifiers import AegisClassifier
+from nemo_curator.datasets import DocumentDataset
+from nemo_curator.utils.distributed_utils import get_client
+from nemo_curator.utils.script_utils import ArgumentHelper
+
+
+def main(args):
+    global_st = time.time()
+
+    # Input can be a string or list
+    input_file_path = "/path/to/data"
+    output_file_path = "./"
+    huggingface_token = "hf_1234"  # Replace with a HuggingFace user access token
+
+    client_args = ArgumentHelper.parse_client_args(args)
+    client_args["cluster_type"] = "gpu"
+    client = get_client(**client_args)
+
+    input_dataset = DocumentDataset.read_json(
+        input_file_path, backend="cudf", add_filename=True
+    )
+
+    safety_classifier = AegisClassifier(
+        aegis_variant="nvidia/Aegis-AI-Content-Safety-LlamaGuard-Permissive-1.0",
+        token=huggingface_token,
+        filter_by=["safe", "O13"],
+    )
+    result_dataset = safety_classifier(dataset=input_dataset)
+
+    result_dataset.to_json(output_file_dir=output_file_path, write_to_filename=True)
+
+    global_et = time.time()
+    print(
+        f"Total time taken for AEGIS classifier inference: {global_et-global_st} s",
+        flush=True,
+    )
+
+    client.close()
+
+
+def attach_args(
+    parser=argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    ),
+):
+    argumentHelper = ArgumentHelper(parser)
+    argumentHelper.add_distributed_classifier_cluster_args()
+
+    return argumentHelper.parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(attach_args().parse_args())

examples/domain_classifier_example.py → examples/classifiers/domain_example.py

@@ -15,7 +15,7 @@
 import argparse
 import time
 
-from nemo_curator import DomainClassifier
+from nemo_curator.classifiers import DomainClassifier
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.utils.distributed_utils import get_client
 from nemo_curator.utils.script_utils import ArgumentHelper
@@ -28,7 +28,9 @@ def main(args):
     input_file_path = "/path/to/data"
     output_file_path = "./"
 
-    client = get_client(**ArgumentHelper.parse_client_args(args))
+    client_args = ArgumentHelper.parse_client_args(args)
+    client_args["cluster_type"] = "gpu"
+    client = get_client(**client_args)
 
     input_dataset = DocumentDataset.read_json(
         input_file_path, backend="cudf", add_filename=True
@@ -54,17 +56,9 @@ def attach_args(
     ),
 ):
     argumentHelper = ArgumentHelper(parser)
+    argumentHelper.add_distributed_classifier_cluster_args()
 
-    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
+    return argumentHelper.parser.parse_args()
 
 
 if __name__ == "__main__":

examples/quality_classifier_example.py → examples/classifiers/quality_example.py

@@ -15,7 +15,7 @@
 import argparse
 import time
 
-from nemo_curator import QualityClassifier
+from nemo_curator.classifiers import QualityClassifier
 from nemo_curator.datasets import DocumentDataset
 from nemo_curator.utils.distributed_utils import get_client
 from nemo_curator.utils.script_utils import ArgumentHelper
@@ -30,7 +30,9 @@ def main(args):
     input_file_path = "/path/to/data"
     output_file_path = "./"
 
-    client = get_client(**ArgumentHelper.parse_client_args(args))
+    client_args = ArgumentHelper.parse_client_args(args)
+    client_args["cluster_type"] = "gpu"
+    client = get_client(**client_args)
 
     input_dataset = DocumentDataset.read_json(
         input_file_path, backend="cudf", add_filename=True
@@ -41,7 +43,8 @@ def main(args):
         filter_by=["High", "Medium"],
     )
     result_dataset = quality_classifier(dataset=input_dataset)
-    print(result_dataset.df.head())
+
+    result_dataset.to_json(output_file_dir=output_file_path, write_to_filename=True)
 
     global_et = time.time()
     print(
@@ -58,17 +61,9 @@ def attach_args(
     ),
 ):
     argumentHelper = ArgumentHelper(parser)
+    argumentHelper.add_distributed_classifier_cluster_args()
 
-    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
+    return argumentHelper.parser.parse_args()
 
 
 if __name__ == "__main__":

nemo_curator/classifiers/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+os.environ["RAPIDS_NO_INITIALIZE"] = "1"
+from .aegis import AegisClassifier
+from .domain import DomainClassifier
+from .quality import QualityClassifier
+
+__all__ = ["DomainClassifier", "QualityClassifier", "AegisClassifier"]