> 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.models.model

## Module Contents

### Classes

| Name                                                              | Description                                  |
| ----------------------------------------------------------------- | -------------------------------------------- |
| [`ModelStage`](#nemo_curator-stages-text-models-model-ModelStage) | Base class for Hugging Face model inference. |

### API

```python
class nemo_curator.stages.text.models.model.ModelStage(
    model_identifier: str,
    cache_dir: str | None = None,
    hf_token: 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,
    unpack_inference_batch: bool = False,
    autocast: bool = True
)
```

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

Base class for Hugging Face model inference.

**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.

---

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

Hugging Face token for downloading the model, if needed. 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.

---

**`unpack_inference_batch`** `bool` — default: False

Whether to unpack the inference batch with \*\*kwargs. Defaults to False.

---

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

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

---

**`name`**

---

**`resources`** `= Resources(cpus=1, gpus=1)`

---

```python
nemo_curator.stages.text.models.model.ModelStage._model_forward(
    model_input_batch: dict[str, torch.Tensor]
) -> torch.Tensor
```

```python
nemo_curator.stages.text.models.model.ModelStage.collect_outputs(
    processed_outputs: list[dict[str, numpy.ndarray]]
) -> dict[str, numpy.ndarray]
```

```python
nemo_curator.stages.text.models.model.ModelStage.create_output_dataframe(
    df_cpu: pandas.DataFrame,
    collected_output: dict[str, numpy.ndarray]
) -> pandas.DataFrame
```

```python
nemo_curator.stages.text.models.model.ModelStage.inputs() -> tuple[list[str], list[str]]
```

```python
nemo_curator.stages.text.models.model.ModelStage.outputs() -> tuple[list[str], list[str]]
```

```python
nemo_curator.stages.text.models.model.ModelStage.process(
    batch: nemo_curator.tasks.DocumentBatch
) -> nemo_curator.tasks.DocumentBatch
```

```python
nemo_curator.stages.text.models.model.ModelStage.process_model_output(
    outputs: torch.Tensor,
    model_input_batch: dict[str, torch.Tensor] | None = None
) -> dict[str, numpy.ndarray] | torch.Tensor
```

```python
nemo_curator.stages.text.models.model.ModelStage.setup(
    _: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None
```

```python
nemo_curator.stages.text.models.model.ModelStage.setup_on_node(
    _node_info: nemo_curator.backends.base.NodeInfo | None = None,
    _worker_metadata: nemo_curator.backends.base.WorkerMetadata = None
) -> None
```

```python
nemo_curator.stages.text.models.model.ModelStage.teardown() -> None
```

```python
nemo_curator.stages.text.models.model.ModelStage.yield_next_batch(
    df: pandas.DataFrame
) -> collections.abc.Generator[dict[str, torch.Tensor]]
```

Yields a generator of model inputs for the next batch.
We only move the batch to the GPU to reduce the memory overhead.

**Parameters:**

**`df`** `pd.DataFrame`

The Pandas DataFrame (with input\_ids and attention\_mask) to process.

---