Adding Custom Models

Learn how to integrate custom models into NeMo Curator stages.

The NeMo Curator container includes a robust set of default models, but you can add your own for specialized tasks.

Before You Start

Before you begin adding a custom model, make sure that you have:

Reviewed the pipeline concepts and diagrams.
A working NeMo Curator development environment.
Optionally prepared a container image that includes your model dependencies.
Optionally created a custom environment to support your new custom model.

How to Add a Custom Model

Review Model Interface

In NeMo Curator, models inherit from nemo_curator.models.base.ModelInterface and must implement model_id_names and setup:

1 class ModelInterface(abc.ABC):
2     """Abstract base class for models used inside stages."""
3 
4     @property
5     @abc.abstractmethod
6     def model_id_names(self) -> list[str]:
7         """Return a list of model IDs associated with this model (for example, Hugging Face IDs)."""
8 
9     @abc.abstractmethod
10     def setup(self) -> None:
11         """Set up the model (load weights, allocate resources)."""

Create New Model

For this tutorial, we’ll sketch a minimal model for demonstration.

1 from typing import Optional
2 
3 import numpy as np
4 import numpy.typing as npt
5 import torch
6 
7 from nemo_curator.models.base import ModelInterface
8 
9 WEIGHTS_MODEL_ID = "example/my-model"
10 
11 class MyCore(torch.nn.Module):
12     def __init__(self, resolution: int = 224):
13         super().__init__()
14         self.resolution = resolution
15         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
16         # Initialize your network here
17         self.net = torch.nn.Identity().to(self.device)
18 
19     @torch.no_grad()
20     def __call__(self, x: npt.NDArray[np.float32]) -> torch.Tensor:
21         tensor = torch.from_numpy(x).to(self.device).float()
22         return self.net(tensor)
23 
24 class MyModel(ModelInterface):
25     def __init__(self, model_dir: str, resolution: int = 224) -> None:
26         self.model_dir = model_dir
27         self.resolution = resolution
28         self._model: Optional[MyCore] = None
29 
30     def model_id_names(self) -> list[str]:
31         return [WEIGHTS_MODEL_ID]
32 
33     def setup(self) -> None:
34         # Load weights from self.model_dir/WEIGHTS_MODEL_ID if needed
35         self._model = MyCore(self.resolution)
36         self._model.eval()

Let’s go through each part of the code piece by piece.

Define the PyTorch Model

1 WEIGHTS_MODEL_ID = "example/my-model"  # your huggingface (or other) model id
2 
3 class MyCore(torch.nn.Module):
4     def __init__(self, resolution: int = 224):
5         super().__init__()
6         # Initialize network and load weights from a local path derived from model_dir and WEIGHTS_MODEL_ID

Provide a model ID (for example, a HuggingFace ID) if you plan to cache or fetch weights. The pipeline can download weights prior to setup() via your model class method if you provide one.

Implement the Model Interface

1 class MyModel(ModelInterface):
2 	...

Your model implements the interface. It defines methods to declare weight identifiers and to initialize the underlying core network.

1     def setup(self) -> None:
2         self._model = MyCore(self.resolution)
3         self._model.eval()

The setup method initializes the underlying MyCore class that performs the model inference.

1     def model_id_names(self) -> list[str]:
2         return [WEIGHTS_MODEL_ID]

The model_id_names property returns a list of weight IDs. These typically correspond to model repository names but do not have to.

If your stage requires a specific environment, manage that in the stage’s resources (for example, gpu_memory_gb or gpus) and container image, rather than on the model. GPU allocation is managed at the stage level using Resources, not on the model.

Manage model weights

Provide your model with a model_dir where weights are stored. Your stage should ensure that any required weights are available at runtime (for example, by mounting them into the container or downloading them prior to execution).

Next Steps

Now that you have created a custom model, you can create a custom stage that uses your code.