NeMo NLP collection API

Model Classes

Modules

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

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

input_example()[source]

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)[source]

Restores module/model with weights

restore_weights(restore_path: str)[source]

Restores module/model’s weights

classmethod save_to(save_path: str)[source]

Saves module/model with weights

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

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()[source]
forward(input_ids, attention_mask, token_type_ids)[source]
property hidden_size

Property returning hidden size.

Returns

Hidden size.

restore_weights(restore_path: str)[source]
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)[source]

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)[source]
class nemo.collections.nlp.modules.BertEncoder(*args: Any, **kwargs: Any)[source]

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)[source]
class nemo.collections.nlp.modules.DistilBertEncoder(*args: Any, **kwargs: Any)[source]

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)[source]
class nemo.collections.nlp.modules.RobertaEncoder(*args: Any, **kwargs: Any)[source]

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)[source]
class nemo.collections.nlp.modules.SequenceClassifier(*args: Any, **kwargs: Any)[source]

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

forward(hidden_states)[source]
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)[source]

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[source]

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)[source]

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

forward(hidden_states)[source]
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[source]

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][source]

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][source]

Returns the list of supported Megatron-LM models