NeMo NLP collection API

Model Classes

Modules

class nemo.collections.nlp.modules.BertModule(*args: Any, **kwargs: Any)

Bases: torch.nn., nemo.core.classes.common.Typing, nemo.core.classes.common.Serialization, nemo.core.classes.common.FileIO

input_example()

Generates input examples for tracing etc. :returns: A tuple of input examples.

property input_types: Optional[Dict[str, nemo.core.neural_types.neural_type.NeuralType]]

Define these to enable input neural type checks

property output_types: Optional[Dict[str, nemo.core.neural_types.neural_type.NeuralType]]

Define these to enable output neural type checks

classmethod restore_from(restore_path: str)

Restores module/model with weights

restore_weights(restore_path: str)

Restores module/model’s weights

classmethod save_to(save_path: str)

Saves module/model with weights

class nemo.collections.nlp.modules.MegatronBertEncoder(*args: Any, **kwargs: Any)

Bases: torch.nn., nemo.core.classes.common.Typing, nemo.core.classes.common.Serialization, nemo.core.classes.common.FileIO

MegatronBERT wraps around the Megatron Language model from https://github.com/NVIDIA/Megatron-LM

Parameters

• config_file (str) – path to model configuration file.

• vocab_file (str) – path to vocabulary file.

• tokenizer_type (str) – tokenizer type, currently only ‘BertWordPieceLowerCase’ supported.

complete_lazy_init()

forward(input_ids, attention_mask, token_type_ids)

property hidden_size

Property returning hidden size.

Returns

Hidden size.

restore_weights(restore_path: str)

Restores module/model’s weights.

For model parallel checkpoints the directory structure should be restore_path/mp_rank_0X/model_optim_rng.pt

Parameters

restore_path (str) – restore_path should a file or a directory if using model parallel

class nemo.collections.nlp.modules.AlbertEncoder(*args: Any, **kwargs: Any)

Bases: transformers., nemo.collections.nlp.modules.common.bert_module.BertModule

Wraps around the Huggingface transformers implementation repository for easy use within NeMo.

forward(input_ids, attention_mask, token_type_ids)

class nemo.collections.nlp.modules.BertEncoder(*args: Any, **kwargs: Any)

Bases: transformers., nemo.collections.nlp.modules.common.bert_module.BertModule

Wraps around the Huggingface transformers implementation repository for easy use within NeMo.

forward(input_ids, attention_mask, token_type_ids)

class nemo.collections.nlp.modules.DistilBertEncoder(*args: Any, **kwargs: Any)

Bases: transformers., nemo.collections.nlp.modules.common.bert_module.BertModule

Wraps around the Huggingface transformers implementation repository for easy use within NeMo.

forward(input_ids, attention_mask, token_type_ids=None)

class nemo.collections.nlp.modules.RobertaEncoder(*args: Any, **kwargs: Any)

Bases: transformers., nemo.collections.nlp.modules.common.bert_module.BertModule

Wraps around the Huggingface transformers implementation repository for easy use within NeMo.

forward(input_ids, token_type_ids, attention_mask)

class nemo.collections.nlp.modules.SequenceClassifier(*args: Any, **kwargs: Any)

Bases: torch.nn., nemo.core.classes.common.Typing, nemo.core.classes.common.Serialization, nemo.core.classes.common.FileIO

forward(hidden_states)

property output_types: Optional[Dict[str, nemo.core.neural_types.neural_type.NeuralType]]

Define these to enable output neural type checks

class nemo.collections.nlp.modules.SequenceRegression(*args: Any, **kwargs: Any)

Bases: torch.nn., nemo.core.classes.common.Typing, nemo.core.classes.common.Serialization, nemo.core.classes.common.FileIO

Parameters

• hidden_size – the hidden size of the mlp head on the top of the encoder

• num_layers – number of the linear layers of the mlp head on the top of the encoder

• activation – type of activations between layers of the mlp head

• dropout – the dropout used for the mlp head

• use_transformer_init – initializes the weights with the same approach used in Transformer

• idx_conditioned_on – index of the token to use as the sequence representation for the classification task, default is the first token

forward(hidden_states: torch.Tensor) → torch.Tensor

Forward pass through the module.

Parameters

hidden_states – hidden states for each token in a sequence, for example, BERT module output

property output_types: Optional[Dict[str, nemo.core.neural_types.neural_type.NeuralType]]

Define these to enable output neural type checks

class nemo.collections.nlp.modules.SequenceTokenClassifier(*args: Any, **kwargs: Any)

Bases: torch.nn., nemo.core.classes.common.Typing, nemo.core.classes.common.Serialization, nemo.core.classes.common.FileIO

forward(hidden_states)

property output_types: Optional[Dict[str, nemo.core.neural_types.neural_type.NeuralType]]

Define these to enable output neural type checks

nemo.collections.nlp.modules.get_lm_model(pretrained_model_name: str, config_dict: Optional[dict] = None, config_file: Optional[str] = None, checkpoint_file: Optional[str] = None, vocab_file: Optional[str] = None) → nemo.collections.nlp.modules.common.bert_module.BertModule

Helper function to instantiate a language model encoder, either from scratch or a pretrained model. If only pretrained_model_name are passed, a pretrained model is returned. If a configuration is passed, whether as a file or dictionary, the model is initialized with random weights.

Parameters

• pretrained_model_name – pretrained model name, for example, bert-base-uncased or megatron-bert-cased. See get_pretrained_lm_models_list() for full list.

• config_dict – path to the model configuration dictionary

• config_file – path to the model configuration file

• checkpoint_file – path to the pretrained model checkpoint

• vocab_file – path to vocab_file to be used with Megatron-LM

Returns

Pretrained BertModule

nemo.collections.nlp.modules.get_pretrained_lm_models_list(include_external: bool = False) → List[str]

Returns the list of supported pretrained model names

Parameters

• if true includes all HuggingFace model names (include_external) –

• only those supported language models in NeMo. (not) –

nemo.collections.nlp.modules.get_megatron_lm_models_list() → List[str]

Returns the list of supported Megatron-LM models