Add AEGIS classifier

[ryantwolf]

Description

Adds the AEGIS Classifier and refactors the layout of the distributed data classifiers.

Usage

from nemo_curator.classifiers import AegisClassifier
from nemo_curator.datasets import DocumentDataset
from nemo_curator import get_client

client = get_client(cluster_type="gpu")

input_dataset = DocumentDataset.read_json(
   "input/", backend="cudf", add_filename=True
)

safety_classifier = AegisClassifier(
    aegis_variant="nvidia/Aegis-AI-Content-Safety-LlamaGuard-Permissive-1.0",
    token="hf_1234",
    filter_by=["safe", "O13"],
)
result_dataset = safety_classifier(dataset=input_dataset)

result_dataset.to_json("output/", write_to_filename=True)

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
• The documentation is up to date with these changes.

The docs will be updated in a separate PR for the technical reviewer to review so they don't block this PR.

[ryantwolf]

Should be ready for review now @sarahyurick @VibhuJawa. We need to wait on merging it in though until this PR is merged in crossfit.

[sarahyurick]

Tested with https://github.com/NVIDIA/NeMo-Curator/blob/main/tutorials/distributed_data_classification/distributed_data_classification.ipynb to confirm domain and quality classification are still intact.

I'm currently not able to run the AEGIS classifier, though.

[sarahyurick]

When testing your example in the PR description (and using rapidsai/crossfit#66), I get:

TypeError: HFModel.__init__() got an unexpected keyword argument 'start_batch_size'

?

@VibhuJawa I see start_batch_size is in rapidsai/crossfit#66, do you know what could be causing this error?

[VibhuJawa]

I think you might have installation issues, like if you pip install this crossfit and then this PR , there is a chance we end up using mainline. The fix is just to install the crossfit PR again.

I also ran into that when testing this PR, so that will be my best guess.

[VibhuJawa]

Mostly looks good, have requested some minor changes

[VibhuJawa]

Suggested change

	self.model, cols=["_hidden_text"], tokenizer_type="sentencepiece"
	self.model, cols=["_hidden_text"], tokenizer_type="default"

[ryantwolf]

I made the change, but what is this doing exactly?

[VibhuJawa]

So there are mostly 3 flavors of tokenizers in use today:

• subword : Used by bert like models
• BPE (Byte Pair Encoding) : Used by GPT like models
• sentence-piece: Another flavor of BPE, Used by models like llama-3 etc.

Today, only subword is GPU accelerated, in the future BPE and SentencePiece will be GPU accelerated but in future we shall accelerate all of them

https://github.com/rapidsai/crossfit/blob/97febffff4bbc32317afaae0aaa94b57d13645c6/crossfit/op/tokenize.py#L31-L34

What default does is (which is functionally same as sentence-piece) is use the provided tokenizer in the model directly. See below:

https://github.com/rapidsai/crossfit/blob/97febffff4bbc32317afaae0aaa94b57d13645c6/crossfit/op/tokenize.py#L59-L75

Why it matters:
Even though there is no functionality difference b/w them today, it might be there in the future.
Having detault means that no unexpected changes happen here when we update crossfit.

[ryantwolf]

@sarahyurick @VibhuJawa thanks for the reviews. Should be good for another round.

[VibhuJawa]

Mostly there, one small nit. Thanks for pushing on this @ryantwolf

[VibhuJawa]

LGTM to me. Thanks again @ryantwolf

[sarahyurick]

Can confirm domain and quality classifiers work on my end. I'm still working on running the AEGIS classifier, though.

[sarahyurick]

LGTM, thanks!