> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/curator/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/curator/_mcp/server.

# nemo_curator.stages.text.classifiers.base

## Module Contents

### Classes

| Name                                                                                                | Description                                               |
| --------------------------------------------------------------------------------------------------- | --------------------------------------------------------- |
| [`ClassifierModelStage`](#nemo_curator-stages-text-classifiers-base-ClassifierModelStage)           | Stage for Hugging Face model inference.                   |
| [`Deberta`](#nemo_curator-stages-text-classifiers-base-Deberta)                                     | Base PyTorch model where we add a classification head.    |
| [`DistributedDataClassifier`](#nemo_curator-stages-text-classifiers-base-DistributedDataClassifier) | Base composite stage for distributed data classification. |

### API

```python
class nemo_curator.stages.text.classifiers.base.ClassifierModelStage(
    model_identifier: str,
    cache_dir: str | None = None,
    label_field: str = 'preds',
    score_field: str | None = None,
    model_inference_batch_size: int = 256,
    has_seq_order: bool = True,
    padding_side: typing.Literal['left', 'right'] = 'right',
    max_seq_length: int | None = None,
    autocast: bool = True,
    keep_tokens: bool = False
)
```

**Bases:** [ModelStage](/nemo-curator/nemo_curator/stages/text/models/model#nemo_curator-stages-text-models-model-ModelStage)

Stage for Hugging Face model inference.

**Parameters:**

**`model_identifier`** `str`

The identifier of the Hugging Face model.

---

**`label_field`** `str` — default: 'preds'

The name of the prediction column.

---

**`score_field`** `str | None` — default: None

The name of the probability column. Defaults to None.

---

**`model_inference_batch_size`** `int` — default: 256

The size of the batch for model inference. Defaults to 256.

---

**`has_seq_order`** `bool` — default: True

Whether to sort the input data by the length of the input tokens.
Sorting is encouraged to improve the performance of the inference model. Defaults to True.

---

**`padding_side`** `Literal['left', 'right']` — default: 'right'

The side to pad the input tokens. Defaults to "right".

---

**`max_seq_length`** `int | None` — default: None

If provided, clips the input tokens before the forward pass. Defaults to None.

---

**`autocast`** `bool` — default: True

Whether to use autocast. When True, we trade off minor accuracy for faster inference.
Defaults to True.

---

**`keep_tokens`** `bool` — default: False

Whether to keep the input tokens in the output dataframe. Defaults to False.

---

```python
nemo_curator.stages.text.classifiers.base.ClassifierModelStage._setup(
    local_files_only: bool = True
) -> None
```

```python
nemo_curator.stages.text.classifiers.base.ClassifierModelStage.create_output_dataframe(
    df_cpu: pandas.DataFrame,
    collected_output: dict[str, numpy.ndarray]
) -> pandas.DataFrame
```

```python
nemo_curator.stages.text.classifiers.base.ClassifierModelStage.outputs() -> tuple[list[str], list[str]]
```

```python
nemo_curator.stages.text.classifiers.base.ClassifierModelStage.process_model_output(
    outputs: torch.Tensor,
    _: dict[str, torch.Tensor] | None = None
) -> dict[str, numpy.ndarray]
```

```python
class nemo_curator.stages.text.classifiers.base.Deberta(
    config: dataclasses.dataclass
)
```

**Bases:** `Module`, `PyTorchModelHubMixin`

Base PyTorch model where we add a classification head.

**Parameters:**

**`config`** `dataclass`

The configuration of the model.

---

**`device`** `device`

---

**`dropout`** `= nn.Dropout(config['fc_dropout'])`

---

**`fc`**

---

**`model`** `= AutoModel.from_pretrained(config['base_model'])`

---

```python
nemo_curator.stages.text.classifiers.base.Deberta.forward(
    batch: dict[str, torch.Tensor]
) -> torch.Tensor
```

```python
class nemo_curator.stages.text.classifiers.base.DistributedDataClassifier(
    model_identifier: str,
    cache_dir: str | None = None,
    label_field: str = 'preds',
    score_field: str | None = None,
    text_field: str = 'text',
    filter_by: list[str] | None = None,
    max_chars: int | None = None,
    max_seq_length: int | None = None,
    padding_side: typing.Literal['left', 'right'] = 'right',
    sort_by_length: bool = True,
    model_inference_batch_size: int = 256,
    autocast: bool = True,
    keep_tokens: bool = False,
    use_existing_tokens: bool = False
)
```

Dataclass

**Bases:** [CompositeStage\[DocumentBatch, DocumentBatch\]](/nemo-curator/nemo_curator/stages/base#nemo_curator-stages-base-CompositeStage)

Base composite stage for distributed data classification.

It decomposes into a tokenizer stage and a model stage.

**Parameters:**

**`model_identifier`** `str`

The identifier of the Hugging Face model.

---

**`cache_dir`** `str | None` — default: None

The Hugging Face cache directory. Defaults to None.

---

**`label_field`** `str` — default: 'preds'

The name of the prediction column. Defaults to "preds".

---

**`score_field`** `str | None` — default: None

The name of the probability column. Defaults to None.

---

**`text_field`** `str` — default: 'text'

The name of the text field in the input data. Defaults to "text".

---

**`filter_by`** `list[str] | None` — default: None

For categorical classifiers, the list of labels to filter the data by. Defaults to None.

---

**`max_chars`** `int | None` — default: None

Limits the total number of characters that can be fed to the tokenizer.
If None, text will not be truncated. Defaults to None.

---

**`max_seq_length`** `int | None` — default: None

Limits the total sequence returned by the tokenizer so that it has a maximum length.
If None, the tokenizer's model\_max\_length is used. Defaults to 512.

---

**`padding_side`** `Literal['left', 'right']` — default: 'right'

The side to pad the input tokens. Defaults to "right".

---

**`sort_by_length`** `bool` — default: True

Whether to sort the input data by the length of the input tokens.
Sorting is encouraged to improve the performance of the inference model. Defaults to True.

---

**`model_inference_batch_size`** `int` — default: 256

The size of the batch for model inference. Defaults to 256.

---

**`autocast`** `bool` — default: True

Whether to use autocast. When True, we trade off minor accuracy for faster inference.
Defaults to True.

---

**`keep_tokens`** `bool` — default: False

Whether to keep the input tokens in the output dataframe. Defaults to False.

---

**`use_existing_tokens`** `bool` — default: False

Whether to use the existing tokens from the input dataframe.
If True, assume the relevant token fields are \["input\_ids", "attention\_mask"] and skip tokenization.
Defaults to False.

---

**`autocast`** `bool = True`

---

**`cache_dir`** `str | None = None`

---

**`filter_by`** `list[str] | None = None`

---

**`keep_tokens`** `bool = False`

---

**`label_field`** `str = 'preds'`

---

**`max_chars`** `int | None = None`

---

**`max_seq_length`** `int | None = None`

---

**`model_identifier`** `str`

---

**`model_inference_batch_size`** `int = 256`

---

**`padding_side`** `Literal['left', 'right'] = 'right'`

---

**`score_field`** `str | None = None`

---

**`sort_by_length`** `bool = True`

---

**`text_field`** `str = 'text'`

---

**`use_existing_tokens`** `bool = False`

---

```python
nemo_curator.stages.text.classifiers.base.DistributedDataClassifier.__post_init__() -> None
```

```python
nemo_curator.stages.text.classifiers.base.DistributedDataClassifier.decompose() -> list[nemo_curator.stages.base.ProcessingStage]
```

```python
nemo_curator.stages.text.classifiers.base.DistributedDataClassifier.filter_by_category(
    value: str
) -> bool
```

```python
nemo_curator.stages.text.classifiers.base.DistributedDataClassifier.inputs() -> tuple[list[str], list[str]]
```

```python
nemo_curator.stages.text.classifiers.base.DistributedDataClassifier.outputs() -> tuple[list[str], list[str]]
```