nemo_curator.models.base

View as Markdown

Module Contents

Classes

NameDescription
ModelInterfaceAbstract base class that defines an interface for machine learning models.

API

class nemo_curator.models.base.ModelInterface()
Abstract

Bases: ABC

Abstract base class that defines an interface for machine learning models.

Specifically focused on their weight handling and environmental setup.

This interface allows our pipeline code to download weights locally and setup models in a uniform way. It does not place any restrictions on how inference is run.

model_id_names
list[str]

Returns a list of model IDs associated with the model.

In cosmos-curate, each model has an ID associated with it. This is often the huggingspace name for that model (e.g. Salesforce/instructblip-vicuna-13b).

nemo_curator.models.base.ModelInterface.setup() -> None
abstract

Set up the model for use, such as loading weights and building computation graphs.