nemo_curator.models.nsfw

Module Contents

Classes

Name	Description
`NSFWModel`	NSFW detection model.
`NSFWScorer`	Public interface for NSFW scoring of image embeddings.
`Normalization`	Normalization layer for NSFW model.

Data

_NSFW_MODEL_ID

API

class nemo_curator.models.nsfw.NSFWModel()

Bases: Module

NSFW detection model.

A neural network that processes CLIP embeddings to predict NSFW scores. Based on LAION’s CLIP-based-NSFW-Detector.

act

= nn.ReLU()

act_out

= nn.Sigmoid()

linear_1

= nn.Linear(768, 64)

linear_2

= nn.Linear(64, 512)

linear_3

= nn.Linear(512, 256)

linear_4

= nn.Linear(256, 1)

norm

= Normalization([768])

nemo_curator.models.nsfw.NSFWModel.forward(
    x: torch.Tensor
) -> torch.Tensor

Forward pass through the NSFW model.

Parameters:

x

torch.Tensor

Input embeddings tensor.

Returns: torch.Tensor

NSFW probability scores.

class nemo_curator.models.nsfw.NSFWScorer(
    model_dir: str
)

Bases: ModelInterface

Public interface for NSFW scoring of image embeddings.

This class provides a standardized interface for scoring the likelihood of images containing sexually explicit material using a pre-trained model.

conda_env_name

str

Get the name of the conda environment required for this model.

device

= 'cuda' if torch.cuda.is_available() else 'cpu'

dtype

= torch.float32

model_id_names

list[str]

Get the model ID names associated with this NSFW scorer.

nemo_curator.models.nsfw.NSFWScorer.__call__(
    embeddings: torch.Tensor | numpy.typing.NDArray[numpy.float32]
) -> torch.Tensor

Score the NSFW likelihood of input embeddings.

Parameters:

embeddings

torch.Tensor | npt.NDArray[np.float32]

Input embeddings as either a torch tensor or numpy array.

Returns: torch.Tensor

NSFW probability scores for each input embedding.

nemo_curator.models.nsfw.NSFWScorer.download_weights_on_node(
    model_dir: str
) -> None

classmethod

Download NSFW model weights from LAION repository.

Parameters:

model_dir

str

Directory to download the weights to.

nemo_curator.models.nsfw.NSFWScorer.setup() -> None

Set up the NSFW scoring model by loading weights.

class nemo_curator.models.nsfw.Normalization(
    shape: list[int]
)

Bases: Module

Normalization layer for NSFW model.

Applies normalization to input tensors using pre-computed mean and variance.

nemo_curator.models.nsfw.Normalization.forward(
    x: torch.Tensor
) -> torch.Tensor

Apply normalization to input tensor.

Parameters:

x

torch.Tensor

Input tensor to normalize.

Returns: torch.Tensor

Normalized tensor.

nemo_curator.models.nsfw._NSFW_MODEL_ID = 'laion/clip-autokeras-binary-nsfw'

nemo_curator.models.nsfw._URL_MAPPING = {'laion/clip-autokeras-binary-nsfw': 'https://github.com/LAION-AI/CLIP-based-NSF...