NeMo ASR collection API#
Model Classes#
- class nemo.collections.asr.models.EncDecCTCModel(*args: Any, **kwargs: Any)[source]#
Bases:
ASRModel
,ExportableEncDecModel
,ASRModuleMixin
,InterCTCMixin
Base class for encoder decoder CTC-based models.
- change_decoding_strategy(decoding_cfg: omegaconf.DictConfig)[source]#
Changes decoding strategy used during CTC decoding process.
- Parameters
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
- change_vocabulary(new_vocabulary: List[str], decoding_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Changes vocabulary used during CTC decoding process. Use this method when fine-tuning on from pre-trained model. This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would use it if you want to use pretrained encoder when fine-tuning on a data in another language, or when you’d need model to learn capitalization, punctuation and/or special characters.
If new_vocabulary == self.decoder.vocabulary then nothing will be changed.
- Parameters
new_vocabulary – list with new vocabulary. Must contain at least 2 elements. Typically, this is target alphabet.
Returns: None
- forward(input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None)[source]#
Forward pass of the model.
- Parameters
input_signal – Tensor that represents a batch of raw audio signals, of shape [B, T]. T here represents timesteps, with 1 second of audio represented as self.sample_rate number of floating point values.
input_signal_length – Vector of length B, that contains the individual lengths of the audio sequences.
processed_signal – Tensor that represents a batch of processed audio signals, of shape (B, D, T) that has undergone processing via some DALI preprocessor.
processed_signal_length – Vector of length B, that contains the individual lengths of the processed audio sequences.
- Returns
A tuple of 3 elements - 1) The log probabilities tensor of shape [B, T, D]. 2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B]. 3) The greedy token predictions of the model of shape [B, T] (via argmax)
- property input_types: Optional[Dict[str, NeuralType]]#
Define these to enable input neural type checks
- classmethod list_available_models() List[PretrainedModelInfo] [source]#
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA’s NGC cloud.
- Returns
List of available pre-trained models.
- multi_test_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple test datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- multi_validation_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple validation datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- property output_types: Optional[Dict[str, NeuralType]]#
Define these to enable output neural type checks
- setup_test_data(test_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the test data loader via a Dict-like object.
- Parameters
test_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- setup_training_data(train_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the training data loader via a Dict-like object.
- Parameters
train_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- setup_validation_data(val_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the validation data loader via a Dict-like object.
- Parameters
val_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- transcribe(paths2audio_files: List[str], batch_size: int = 4, logprobs: bool = False, return_hypotheses: bool = False, num_workers: int = 0, channel_selector: Optional[Union[int, Iterable[int], str]] = None, augmentor: Optional[omegaconf.DictConfig] = None, verbose: bool = True) List[str] #
If modify this function, please remember update transcribe_partial_audio() in nemo/collections/asr/parts/utils/trancribe_utils.py
Uses greedy decoding to transcribe audio files. Use this method for debugging and prototyping.
- Parameters
paths2audio_files – (a list) of paths to audio files. Recommended length per file is between 5 and 25 seconds. But it is possible to pass a few hours long file if enough GPU memory is available.
batch_size – (int) batch size to use during inference. Bigger will result in better throughput performance but would use more memory.
logprobs – (bool) pass True to get log probabilities instead of transcripts.
return_hypotheses – (bool) Either return hypotheses or text With hypotheses can do some postprocessing like getting timestamp or rescoring
num_workers – (int) number of workers for DataLoader
channel_selector (int | Iterable[int] | str) – select a single channel or a subset of channels from multi-channel audio. If set to ‘average’, it performs averaging across channels. Disabled if set to None. Defaults to None.
augmentor – (DictConfig): Augment audio samples during transcription if augmentor is applied.
verbose – (bool) whether to display tqdm progress bar
- Returns
A list of transcriptions (or raw log probabilities if logprobs is True) in the same order as paths2audio_files
- class nemo.collections.asr.models.EncDecCTCModelBPE(*args: Any, **kwargs: Any)[source]#
Bases:
EncDecCTCModel
,ASRBPEMixin
Encoder decoder CTC-based models with Byte Pair Encoding.
- change_decoding_strategy(decoding_cfg: omegaconf.DictConfig)[source]#
Changes decoding strategy used during CTC decoding process.
- Parameters
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
- change_vocabulary(new_tokenizer_dir: Union[str, omegaconf.DictConfig], new_tokenizer_type: str, decoding_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Changes vocabulary of the tokenizer used during CTC decoding process. Use this method when fine-tuning on from pre-trained model. This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would use it if you want to use pretrained encoder when fine-tuning on a data in another language, or when you’d need model to learn capitalization, punctuation and/or special characters.
- Parameters
new_tokenizer_dir – Directory path to tokenizer or a config for a new tokenizer (if the tokenizer type is agg)
new_tokenizer_type – Either agg, bpe or wpe. bpe is used for SentencePiece tokenizers, whereas wpe is used for BertTokenizer.
new_tokenizer_cfg – A config for the new tokenizer. if provided, pre-empts the dir and type
Returns: None
- class nemo.collections.asr.models.EncDecRNNTModel(*args: Any, **kwargs: Any)[source]#
Bases:
ASRModel
,ASRModuleMixin
,ExportableEncDecModel
Base class for encoder decoder RNNT-based models.
- change_decoding_strategy(decoding_cfg: omegaconf.DictConfig)[source]#
Changes decoding strategy used during RNNT decoding process.
- Parameters
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
- change_vocabulary(new_vocabulary: List[str], decoding_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning a pre-trained model. This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would use it if you want to use pretrained encoder when fine-tuning on data in another language, or when you’d need model to learn capitalization, punctuation and/or special characters.
- Parameters
new_vocabulary – list with new vocabulary. Must contain at least 2 elements. Typically, this is target alphabet.
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
Returns: None
- property decoder_joint#
- extract_rnnt_loss_cfg(cfg: Optional[omegaconf.DictConfig])[source]#
Helper method to extract the rnnt loss name, and potentially its kwargs to be passed.
- Parameters
cfg –
Should contain loss_name as a string which is resolved to a RNNT loss name. If the default should be used, then default can be used. Optionally, one can pass additional kwargs to the loss function. The subdict should have a keyname as follows : {loss_name}_kwargs.
Note that whichever loss_name is selected, that corresponding kwargs will be selected. For the “default” case, the “{resolved_default}_kwargs” will be used.
Examples
loss_name: "default" warprnnt_numba_kwargs: kwargs2: some_other_val
- Returns
A tuple, the resolved loss name as well as its kwargs (if found).
- forward(input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None)[source]#
Forward pass of the model. Note that for RNNT Models, the forward pass of the model is a 3 step process, and this method only performs the first step - forward of the acoustic model.
Please refer to the training_step in order to see the full forward step for training - which performs the forward of the acoustic model, the prediction network and then the joint network. Finally, it computes the loss and possibly compute the detokenized text via the decoding step.
Please refer to the validation_step in order to see the full forward step for inference - which performs the forward of the acoustic model, the prediction network and then the joint network. Finally, it computes the decoded tokens via the decoding step and possibly compute the batch metrics.
- Parameters
input_signal – Tensor that represents a batch of raw audio signals, of shape [B, T]. T here represents timesteps, with 1 second of audio represented as self.sample_rate number of floating point values.
input_signal_length – Vector of length B, that contains the individual lengths of the audio sequences.
processed_signal – Tensor that represents a batch of processed audio signals, of shape (B, D, T) that has undergone processing via some DALI preprocessor.
processed_signal_length – Vector of length B, that contains the individual lengths of the processed audio sequences.
- Returns
A tuple of 2 elements - 1) The log probabilities tensor of shape [B, T, D]. 2) The lengths of the acoustic sequence after propagation through the encoder, of shape [B].
- property input_types: Optional[Dict[str, NeuralType]]#
Define these to enable input neural type checks
- classmethod list_available_models() List[PretrainedModelInfo] [source]#
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA’s NGC cloud.
- Returns
List of available pre-trained models.
- list_export_subnets()[source]#
Returns default set of subnet names exported for this model First goes the one receiving input (input_example)
- multi_test_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple test datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- multi_validation_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple validation datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- property output_types: Optional[Dict[str, NeuralType]]#
Define these to enable output neural type checks
- setup_optim_normalization()[source]#
Helper method to setup normalization of certain parts of the model prior to the optimization step.
Supported pre-optimization normalizations are as follows:
# Variation Noise injection model: variational_noise: std: 0.0 start_step: 0 # Joint - Length normalization model: normalize_joint_txu: false # Encoder Network - gradient normalization model: normalize_encoder_norm: false # Decoder / Prediction Network - gradient normalization model: normalize_decoder_norm: false # Joint - gradient normalization model: normalize_joint_norm: false
- setup_test_data(test_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the test data loader via a Dict-like object.
- Parameters
test_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- setup_training_data(train_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the training data loader via a Dict-like object.
- Parameters
train_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- setup_validation_data(val_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Sets up the validation data loader via a Dict-like object.
- Parameters
val_data_config – A config that contains the information regarding construction of an ASR Training dataset.
- Supported Datasets:
AudioToCharDALIDataset
- transcribe(paths2audio_files: List[str], batch_size: int = 4, return_hypotheses: bool = False, partial_hypothesis: Optional[List[Hypothesis]] = None, num_workers: int = 0, channel_selector: Optional[Union[int, Iterable[int], str]] = None, augmentor: omegaconf.DictConfig = None, verbose: bool = True) Tuple[List[str], Optional[List[Hypothesis]]] #
Uses greedy decoding to transcribe audio files. Use this method for debugging and prototyping.
- Parameters
paths2audio_files – (a list) of paths to audio files. Recommended length per file is between 5 and 25 seconds. But it is possible to pass a few hours long file if enough GPU memory is available.
batch_size – (int) batch size to use during inference. Bigger will result in better throughput performance but would use more memory.
return_hypotheses – (bool) Either return hypotheses or text
- With hypotheses can do some postprocessing like getting timestamp or rescoring
num_workers: (int) number of workers for DataLoader channel_selector (int | Iterable[int] | str): select a single channel or a subset of channels from multi-channel audio. If set to ‘average’, it performs averaging across channels. Disabled if set to None. Defaults to None. Uses zero-based indexing. augmentor: (DictConfig): Augment audio samples during transcription if augmentor is applied. verbose: (bool) whether to display tqdm progress bar
- Returns
Returns a tuple of 2 items - * A list of greedy transcript texts / Hypothesis * An optional list of beam search transcript texts / Hypothesis / NBestHypothesis.
- class nemo.collections.asr.models.EncDecRNNTBPEModel(*args: Any, **kwargs: Any)[source]#
Bases:
EncDecRNNTModel
,ASRBPEMixin
Base class for encoder decoder RNNT-based models with subword tokenization.
- change_decoding_strategy(decoding_cfg: omegaconf.DictConfig)[source]#
Changes decoding strategy used during RNNT decoding process.
- Parameters
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
- change_vocabulary(new_tokenizer_dir: Union[str, omegaconf.DictConfig], new_tokenizer_type: str, decoding_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Changes vocabulary used during RNNT decoding process. Use this method when fine-tuning on from pre-trained model. This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would use it if you want to use pretrained encoder when fine-tuning on data in another language, or when you’d need model to learn capitalization, punctuation and/or special characters.
- Parameters
new_tokenizer_dir – Directory path to tokenizer or a config for a new tokenizer (if the tokenizer type is agg)
new_tokenizer_type – Type of tokenizer. Can be either agg, bpe or wpe.
decoding_cfg – A config for the decoder, which is optional. If the decoding type needs to be changed (from say Greedy to Beam decoding etc), the config can be passed here.
Returns: None
- class nemo.collections.asr.models.EncDecClassificationModel(*args: Any, **kwargs: Any)[source]#
Bases:
_EncDecBaseModel
Encoder decoder Classification models.
- change_labels(new_labels: List[str])[source]#
Changes labels used by the decoder model. Use this method when fine-tuning on from pre-trained model. This method changes only decoder and leaves encoder and pre-processing modules unchanged. For example, you would use it if you want to use pretrained encoder when fine-tuning on a data in another dataset.
If new_labels == self.decoder.vocabulary then nothing will be changed.
- Parameters
new_labels – list with new labels. Must contain at least 2 elements. Typically, this is set of labels for the dataset.
Returns: None
- forward(input_signal=None, input_signal_length=None, processed_signal=None, processed_signal_length=None)[source]#
- classmethod list_available_models() Optional[List[PretrainedModelInfo]] [source]#
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA’s NGC cloud.
- Returns
List of available pre-trained models.
- multi_test_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple test datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- multi_validation_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple validation datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- property output_types: Optional[Dict[str, NeuralType]]#
Define these to enable output neural type checks
- class nemo.collections.asr.models.EncDecSpeakerLabelModel(*args: Any, **kwargs: Any)[source]#
Bases:
ModelPT
,ExportableEncDecModel
Encoder decoder class for speaker label models. Model class creates training, validation methods for setting up data performing model forward pass. Expects config dict for
preprocessor
Jasper/Quartznet Encoder
Speaker Decoder
- batch_inference(manifest_filepath, batch_size=32, sample_rate=16000, device='cuda')#
Perform batch inference on EncDecSpeakerLabelModel. To perform inference on single audio file, once can use infer_model, get_label or get_embedding
- To map predicted labels, one can do
arg_values = logits.argmax(axis=1) pred_labels = list(map(lambda t : trained_labels[t], arg_values))
- Parameters
manifest_filepath – Path to manifest file
batch_size – batch size to perform batch inference
sample_rate – sample rate of audio files in manifest file
device – compute device to perform operations.
- Returns
The variables below all follow the audio file order in the manifest file. embs: embeddings of files provided in manifest file logits: logits of final layer of EncDecSpeakerLabel Model gt_labels: labels from manifest file (needed for speaker enrollment and testing) trained_labels: Classification labels sorted in the order that they are mapped by the trained model
- forward_for_export(processed_signal, processed_signal_len)[source]#
This forward is used when we need to export the model to ONNX format. Inputs cache_last_channel and cache_last_time are needed to be passed for exporting streaming models. :param input: Tensor that represents a batch of raw audio signals,
of shape [B, T]. T here represents timesteps.
- Parameters
length – Vector of length B, that contains the individual lengths of the audio sequences.
cache_last_channel – Tensor of shape [N, B, T, H] which contains the cache for last channel layers
cache_last_time – Tensor of shape [N, B, H, T] which contains the cache for last time layers N is the number of such layers which need caching, B is batch size, H is the hidden size of activations, and T is the length of the cache
- Returns
the output of the model
- get_embedding(path2audio_file)[source]#
Returns the speaker embeddings for a provided audio file.
- Parameters
path2audio_file – path to an audio wav file
- Returns
speaker embeddings (Audio representations)
- Return type
emb
- get_label(path2audio_file)[source]#
Returns label of path2audio_file from classes the model was trained on. :param path2audio_file: path to audio wav file
- Returns
label corresponding to the trained model
- Return type
label
- infer_file(path2audio_file)#
- Parameters
path2audio_file – path to an audio wav file
- Returns
speaker embeddings (Audio representations) logits: logits corresponding of final layer
- Return type
emb
- property input_types: Optional[Dict[str, NeuralType]]#
Define these to enable input neural type checks
- classmethod list_available_models() List[PretrainedModelInfo] [source]#
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA’s NGC cloud. :returns: List of available pre-trained models.
- multi_test_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple test datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- multi_validation_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Adds support for multiple validation datasets. Should be overriden by subclass, so as to obtain appropriate logs for each of the dataloaders.
- Parameters
outputs – Same as that provided by LightningModule.on_validation_epoch_end() for a single dataloader.
dataloader_idx – int representing the index of the dataloader.
- Returns
A dictionary of values, optionally containing a sub-dict log, such that the values in the log will be pre-pended by the dataloader prefix.
- property output_types: Optional[Dict[str, NeuralType]]#
Define these to enable output neural type checks
- setup_test_data(test_data_layer_params: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
(Optionally) Setups data loader to be used in test
- Parameters
test_data_layer_config – test data layer parameters.
Returns:
- setup_training_data(train_data_layer_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Setups data loader to be used in training
- Parameters
train_data_layer_config – training data layer parameters.
Returns:
- setup_validation_data(val_data_layer_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Setups data loader to be used in validation :param val_data_layer_config: validation data layer parameters.
Returns:
- verify_speakers(path2audio_file1, path2audio_file2, threshold=0.7)#
Verify if two audio files are from the same speaker or not.
- Parameters
path2audio_file1 – path to audio wav file of speaker 1
path2audio_file2 – path to audio wav file of speaker 2
threshold – cosine similarity score used as a threshold to distinguish two embeddings (default = 0.7)
- Returns
True if both audio files are from same speaker, False otherwise
- class nemo.collections.asr.models.hybrid_asr_tts_models.ASRWithTTSModel(*args: Any, **kwargs: Any)[source]#
Bases:
ASRModel
Hybrid ASR-TTS model: a transparent wrapper for ASR model with frozen text-to-spectrogram pretrained model, which allows to use text-only data for training/finetuning Text-only data can be mixed with audio-text pairs
- class ASRModelTypes(value)[source]#
Bases:
PrettyStrEnum
Supported ASR types, needed for training from scratch
- CTC_BPE = 'ctc_bpe'#
- HYBRID_RNNT_CTC_BPE = 'hybrid_rnnt_ctc_bpe'#
- RNNT_BPE = 'rnnt_bpe'#
- asr_model: Union[EncDecRNNTBPEModel, EncDecCTCModelBPE, EncDecHybridRNNTCTCBPEModel]#
- enhancer_model: Optional[SpectrogramEnhancerModel]#
- classmethod from_asr_config(asr_cfg: omegaconf.DictConfig, asr_model_type: Union[str, ASRModelTypes], tts_model_path: Union[str, Path], enhancer_model_path: Optional[Union[Path, str]] = None, trainer: Optional[pytorch_lightning.Trainer] = None)[source]#
Method to construct model from ASR config for training from scratch
- classmethod from_pretrained_models(asr_model_path: Union[str, Path], tts_model_path: Union[str, Path], enhancer_model_path: Optional[Union[Path, str]] = None, asr_model_fuse_bn: bool = False, cfg: Optional[omegaconf.DictConfig] = None, trainer: Optional[pytorch_lightning.Trainer] = None)[source]#
Load model from pretrained ASR and TTS models :param asr_model_path: path to .nemo ASR model checkpoint :param tts_model_path: path to .nemo TTS model checkpoint :param enhancer_model_path: path to .nemo enhancer model checkpoint :param asr_model_fuse_bn: automatically fuse batchnorm layers in ASR model :param cfg: optional config for hybrid model :param trainer: Pytorch-Lightning trainer
- Returns
ASRWithTTSModel instance
- classmethod list_available_models() List[PretrainedModelInfo] [source]#
This method returns a list of pre-trained model which can be instantiated directly from NVIDIA’s NGC cloud. :returns: List of available pre-trained models.
- multi_validation_epoch_end(outputs, dataloader_idx: int = 0)[source]#
Validation epoch end hook for ASR model
- on_test_epoch_end() Optional[Dict[str, Dict[str, torch.Tensor]]] [source]#
Test epoch end hook, forward to ASR model
- on_validation_epoch_end() Optional[Dict[str, Dict[str, torch.Tensor]]] [source]#
Validation epoch end hook, forward to ASR model
- setup_multiple_test_data(test_data_config: Union[omegaconf.DictConfig, Dict])[source]#
Setup multiple test data for ASR Model
- setup_multiple_validation_data(val_data_config: Union[omegaconf.DictConfig, Dict])[source]#
Setup multiple validation data for ASR model
- setup_optimization(optim_config: Optional[Union[omegaconf.DictConfig, Dict]] = None, optim_kwargs: Optional[Dict[str, Any]] = None)[source]#
Setup optimizer and scheduler. Ensure tts model is frozen. Add optimizer and scheduler to asr model, to allow train_step on ASR model
- setup_test_data(test_data_config: Union[omegaconf.DictConfig, Dict])[source]#
Setup test data for ASR model
- setup_training_data(train_data_config: Optional[Union[omegaconf.DictConfig, Dict]])[source]#
Setup training data from config: text-only, audio-text or mixed data.
- setup_validation_data(val_data_config: Union[omegaconf.DictConfig, Dict])[source]#
Setup validation data for ASR model
- training_step(batch: Union[TextOrAudioToTextBatch, TextToTextBatch, DALIOutputs, tuple], batch_nb: int)[source]#
Training step for ASR-TTS model. - construct spectrogram for the batch (from text - using TTS model, from audio - using ASR preprocessor) - call training_step on ASR model
- transcribe(paths2audio_files: List[str], batch_size: int = 4, verbose: bool = True) List[str] [source]#
Transcribe audio data using ASR model
- tts_model: FastPitchModel#
Modules#
- class nemo.collections.asr.modules.ConvASREncoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
,AccessMixin
Convolutional encoder for ASR models. With this class you can implement JasperNet and QuartzNet models.
- Based on these papers:
https://arxiv.org/pdf/1904.03288.pdf https://arxiv.org/pdf/1910.10261.pdf
- input_example(max_batch=1, max_dim=8192)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- class nemo.collections.asr.modules.ConvASRDecoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
,AdapterModuleMixin
Simple ASR Decoder for use with CTC-based models such as JasperNet and QuartzNet
- Based on these papers:
https://arxiv.org/pdf/1904.03288.pdf https://arxiv.org/pdf/1910.10261.pdf https://arxiv.org/pdf/2005.04290.pdf
- add_adapter(name: str, cfg: omegaconf.DictConfig)[source]#
Add an Adapter module to this module.
- Parameters
name – A globally unique name for the adapter. Will be used to access, enable and disable adapters.
cfg – A DictConfig or Dataclass that contains at the bare minimum __target__ to instantiate a new Adapter module.
- input_example(max_batch=1, max_dim=256)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Define these to enable input neural type checks
- property num_classes_with_blank#
- property output_types#
Define these to enable output neural type checks
- property vocabulary#
- class nemo.collections.asr.modules.ConvASRDecoderClassification(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
Simple ASR Decoder for use with classification models such as JasperNet and QuartzNet
- Based on these papers:
- input_example(max_batch=1, max_dim=256)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Define these to enable input neural type checks
- property num_classes#
- property output_types#
Define these to enable output neural type checks
- class nemo.collections.asr.modules.SpeakerDecoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
Speaker Decoder creates the final neural layers that maps from the outputs of Jasper Encoder to the embedding layer followed by speaker based softmax loss. :param feat_in: Number of channels being input to this module :type feat_in: int :param num_classes: Number of unique speakers in dataset :type num_classes: int :param emb_sizes: shapes of intermediate embedding layers (we consider speaker embbeddings from 1st of this layers)
Defaults to [1024,1024]
- Parameters
pool_mode (str) – Pooling strategy type. options are ‘xvector’,’tap’, ‘attention’ Defaults to ‘xvector (mean and variance)’ tap (temporal average pooling: just mean) attention (attention based pooling)
init_mode (str) – Describes how neural network parameters are initialized. Options are [‘xavier_uniform’, ‘xavier_normal’, ‘kaiming_uniform’,’kaiming_normal’]. Defaults to “xavier_uniform”.
- input_example(max_batch=1, max_dim=256)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Define these to enable input neural type checks
- property output_types#
Define these to enable output neural type checks
- class nemo.collections.asr.modules.ConformerEncoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,StreamingEncoder
,Exportable
,AccessMixin
The encoder for ASR model of Conformer. Based on this paper: ‘Conformer: Convolution-augmented Transformer for Speech Recognition’ by Anmol Gulati et al. https://arxiv.org/abs/2005.08100
- Parameters
feat_in (int) – the size of feature channels
n_layers (int) – number of layers of ConformerBlock
d_model (int) – the hidden size of the model
feat_out (int) – the size of the output features Defaults to -1 (means feat_out is d_model)
subsampling (str) – the method of subsampling, choices=[‘vggnet’, ‘striding’, ‘dw-striding’, ‘stacking’, ‘stacking_norm’] Defaults to striding.
subsampling_factor (int) – the subsampling factor which should be power of 2 Defaults to 4.
subsampling_conv_chunking_factor (int) – optionally, force chunk inputs (helpful for large inputs) Should be power of 2, 1 (auto-chunking, default), or -1 (no chunking)
subsampling_conv_channels (int) – the size of the convolutions in the subsampling module Defaults to -1 which would set it to d_model.
reduction (str, Optional) – the method of reduction, choices=[‘pooling’, ‘striding’]. If no value is passed, then no reduction is performed and the models runs with the original 4x subsampling.
reduction_position (int, Optional) – the index of the layer to apply reduction. If -1, apply reduction at the end.
reduction_factor (int) – the reduction factor which should be either 1 or a power of 2 Defaults to 1.
ff_expansion_factor (int) – the expansion factor in feed forward layers Defaults to 4.
self_attention_model (str) –
type of the attention layer and positional encoding ‘rel_pos’: relative positional embedding and Transformer-XL ‘rel_pos_local_attn’: relative positional embedding and Transformer-XL with local attention using
overlapping chunks. Attention context is determined by att_context_size parameter.
’abs_pos’: absolute positional embedding and Transformer Default is rel_pos.
pos_emb_max_len (int) – the maximum length of positional embeddings Defaults to 5000
n_heads (int) – number of heads in multi-headed attention layers Defaults to 4.
att_context_size (List[Union[List[int],int]]) – specifies the context sizes on each side. Each context size should be a list of two integers like [100,100]. A list of context sizes like [[100,100],[100,50]] can also be passed. -1 means unlimited context. Defaults to [-1,-1]
att_context_probs (List[float]) – a list of probabilities of each one of the att_context_size when a list of them is passed. If not specified, uniform distribution is being used. Defaults to None
att_context_style (str) – ‘regular’ or ‘chunked_limited’. Defaults to ‘regular’
xscaling (bool) – enables scaling the inputs to the multi-headed attention layers by sqrt(d_model) Defaults to True.
untie_biases (bool) – whether to not share (untie) the bias weights between layers of Transformer-XL Defaults to True.
conv_kernel_size (int) – the size of the convolutions in the convolutional modules Defaults to 31.
conv_norm_type (str) – the type of the normalization in the convolutional modules Defaults to ‘batch_norm’.
conv_context_size (list) – it can be”causal” or a list of two integers while conv_context_size[0]+conv_context_size[1]+1==conv_kernel_size. None means [(conv_kernel_size-1)//2, (conv_kernel_size-1)//2], and ‘causal’ means [(conv_kernel_size-1), 0]. Defaults to None.
conv_dual_mode (bool) – specifies if convolution should be dual mode when dual_offline mode is being used. When enables, the left half of the convolution kernel would get masked in streaming cases. Defaults to False
dropout (float) – the dropout rate used in all layers except the attention layers Defaults to 0.1.
dropout_pre_encoder (float) – the dropout rate used before the encoder Defaults to 0.1.
dropout_emb (float) – the dropout rate used for the positional embeddings Defaults to 0.1.
dropout_att (float) – the dropout rate used for the attention layer Defaults to 0.0.
stochastic_depth_drop_prob (float) – if non-zero, will randomly drop layers during training. The higher this value, the more often layers are dropped. Defaults to 0.0.
stochastic_depth_mode (str) – can be either “linear” or “uniform”. If set to “uniform”, all layers have the same probability of drop. If set to “linear”, the drop probability grows linearly from 0 for the first layer to the desired value for the final layer. Defaults to “linear”.
stochastic_depth_start_layer (int) – starting layer for stochastic depth. All layers before this will never be dropped. Note that drop probability will be adjusted accordingly if mode is “linear” when start layer is > 1. Defaults to 1.
global_tokens (int) – number of tokens to be used for global attention. Only relevant if self_attention_model is ‘rel_pos_local_attn’. Defaults to 0.
global_tokens_spacing (int) – how far apart the global tokens are Defaults to 1.
global_attn_separate (bool) – whether the q, k, v layers used for global tokens should be separate. Defaults to False.
- change_attention_model(self_attention_model: Optional[str] = None, att_context_size: Optional[List[int]] = None, update_config: bool = True, device: Optional[torch.device] = None)[source]#
Update the self_attention_model which changes the positional encoding and attention layers.
- Parameters
self_attention_model (str) –
type of the attention layer and positional encoding ‘rel_pos’: relative positional embedding and Transformer-XL ‘rel_pos_local_attn’: relative positional embedding and Transformer-XL with local attention using
overlapping windows. Attention context is determined by att_context_size parameter.
’abs_pos’: absolute positional embedding and Transformer If None is provided, the self_attention_model isn’t changed. Defaults to None.
att_context_size (List[int]) – List of 2 ints corresponding to left and right attention context sizes, or None to keep as it is. Defaults to None.
update_config (bool) – Whether to update the config or not with the new attention model. Defaults to True.
device (torch.device) – If provided, new layers will be moved to the device. Defaults to None.
- change_subsampling_conv_chunking_factor(subsampling_conv_chunking_factor: int)[source]#
Update the conv_chunking_factor (int) Default is 1 (auto) Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers
- Parameters
subsampling_conv_chunking_factor (int) –
- property disabled_deployment_input_names#
Implement this method to return a set of input names disabled for export
- property disabled_deployment_output_names#
Implement this method to return a set of output names disabled for export
- forward(audio_signal, length, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None)[source]#
- forward_for_export(audio_signal, length, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None)[source]#
- forward_internal(audio_signal, length, cache_last_channel=None, cache_last_time=None, cache_last_channel_len=None)[source]#
- input_example(max_batch=1, max_dim=256)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property input_types_for_export#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- property output_types_for_export#
Returns definitions of module output ports.
- set_max_audio_length(max_audio_length)[source]#
Sets maximum input length. Pre-calculates internal seq_range mask.
- setup_streaming_params(chunk_size: Optional[int] = None, shift_size: Optional[int] = None, left_chunks: Optional[int] = None, att_context_size: Optional[list] = None, max_context: int = 10000)[source]#
This function sets the needed values and parameters to perform streaming. The configuration would be stored in self.streaming_cfg. The streaming configuration is needed to simulate streaming inference. :param chunk_size: overrides the chunk size :type chunk_size: int :param shift_size: overrides the shift size for chunks :type shift_size: int :param left_chunks: overrides the number of left chunks visible to each chunk :type left_chunks: int :param max_context: the value used for the cache size of last_channel layers if left context is set to infinity (-1)
Defaults to -1 (means feat_out is d_model)
- class nemo.collections.asr.modules.SqueezeformerEncoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
,AccessMixin
The encoder for ASR model of Squeezeformer. Based on this paper: ‘Squeezeformer: An Efficient Transformer for Automatic Speech Recognition’ by Sehoon Kim et al. https://arxiv.org/abs/2206.00888
- Parameters
feat_in (int) – the size of feature channels
n_layers (int) – number of layers of ConformerBlock
d_model (int) – the hidden size of the model
feat_out (int) – the size of the output features Defaults to -1 (means feat_out is d_model)
subsampling (str) – the method of subsampling, choices=[‘vggnet’, ‘striding’, ‘dw_striding’] Defaults to dw_striding.
subsampling_factor (int) – the subsampling factor which should be power of 2 Defaults to 4.
subsampling_conv_channels (int) – the size of the convolutions in the subsampling module Defaults to -1 which would set it to d_model.
ff_expansion_factor (int) – the expansion factor in feed forward layers Defaults to 4.
self_attention_model (str) – type of the attention layer and positional encoding ‘rel_pos’: relative positional embedding and Transformer-XL ‘abs_pos’: absolute positional embedding and Transformer default is rel_pos.
pos_emb_max_len (int) – the maximum length of positional embeddings Defaulst to 5000
n_heads (int) – number of heads in multi-headed attention layers Defaults to 4.
xscaling (bool) – enables scaling the inputs to the multi-headed attention layers by sqrt(d_model) Defaults to True.
untie_biases (bool) – whether to not share (untie) the bias weights between layers of Transformer-XL Defaults to True.
conv_kernel_size (int) – the size of the convolutions in the convolutional modules Defaults to 31.
conv_norm_type (str) – the type of the normalization in the convolutional modules Defaults to ‘batch_norm’.
dropout (float) – the dropout rate used in all layers except the attention layers Defaults to 0.1.
dropout_emb (float) – the dropout rate used for the positional embeddings Defaults to 0.1.
dropout_att (float) – the dropout rate used for the attention layer Defaults to 0.0.
adaptive_scale (bool) – Whether to scale the inputs to each component by affine scale and bias layer. Or use a fixed scale=1 and bias=0.
time_reduce_idx (int) – Optional integer index of a layer where a time reduction operation will occur. All operations beyond this point will only occur at the reduced resolution.
time_recovery_idx (int) – Optional integer index of a layer where the time recovery operation will occur. All operations beyond this point will occur at the original resolution (resolution after primary downsampling). If no value is provided, assumed to be the last layer.
- input_example(max_batch=1, max_dim=256)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- class nemo.collections.asr.modules.RNNEncoder(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
,Exportable
The RNN-based encoder for ASR models. Followed the architecture suggested in the following paper: ‘STREAMING END-TO-END SPEECH RECOGNITION FOR MOBILE DEVICES’ by Yanzhang He et al. https://arxiv.org/pdf/1811.06621.pdf
- Parameters
feat_in (int) – the size of feature channels
n_layers (int) – number of layers of RNN
d_model (int) – the hidden size of the model
proj_size (int) – the size of the output projection after each RNN layer
rnn_type (str) – the type of the RNN layers, choices=[‘lstm, ‘gru’, ‘rnn’]
bidirectional (float) – specifies whether RNN layers should be bidirectional or not Defaults to True.
feat_out (int) – the size of the output features Defaults to -1 (means feat_out is d_model)
subsampling (str) – the method of subsampling, choices=[‘stacking, ‘vggnet’, ‘striding’] Defaults to stacking.
subsampling_factor (int) – the subsampling factor Defaults to 4.
subsampling_conv_channels (int) – the size of the convolutions in the subsampling module for vggnet and striding Defaults to -1 which would set it to d_model.
dropout (float) – the dropout rate used between all layers Defaults to 0.2.
- input_example()[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- class nemo.collections.asr.modules.RNNTDecoder(*args: Any, **kwargs: Any)[source]#
Bases:
AbstractRNNTDecoder
,Exportable
,AdapterModuleMixin
A Recurrent Neural Network Transducer Decoder / Prediction Network (RNN-T Prediction Network). An RNN-T Decoder/Prediction network, comprised of a stateful LSTM model.
- Parameters
prednet –
A dict-like object which contains the following key-value pairs. pred_hidden: int specifying the hidden dimension of the prediction net. pred_rnn_layers: int specifying the number of rnn layers.
Optionally, it may also contain the following: forget_gate_bias: float, set by default to 1.0, which constructs a forget gate
initialized to 1.0. Reference: [An Empirical Exploration of Recurrent Network Architectures](http://proceedings.mlr.press/v37/jozefowicz15.pdf)
- t_max: int value, set to None by default. If an int is specified, performs Chrono Initialization
of the LSTM network, based on the maximum number of timesteps t_max expected during the course of training. Reference: [Can recurrent neural networks warp time?](https://openreview.net/forum?id=SJcKhk-Ab)
- weights_init_scale: Float scale of the weights after initialization. Setting to lower than one
sometimes helps reduce variance between runs.
- hidden_hidden_bias_scale: Float scale for the hidden-to-hidden bias scale. Set to 0.0 for
the default behaviour.
dropout: float, set to 0.0 by default. Optional dropout applied at the end of the final LSTM RNN layer.
vocab_size – int, specifying the vocabulary size of the embedding layer of the Prediction network, excluding the RNNT blank token.
normalization_mode – Can be either None, ‘batch’ or ‘layer’. By default, is set to None. Defines the type of normalization applied to the RNN layer.
random_state_sampling – bool, set to False by default. When set, provides normal-distribution sampled state tensors instead of zero tensors during training. Reference: [Recognizing long-form speech using streaming end-to-end models](https://arxiv.org/abs/1910.11455)
blank_as_pad –
bool, set to True by default. When set, will add a token to the Embedding layer of this prediction network, and will treat this token as a pad token. In essence, the RNNT pad token will be treated as a pad token, and the embedding layer will return a zero tensor for this token.
It is set by default as it enables various batch optimizations required for batched beam search. Therefore, it is not recommended to disable this flag.
- add_adapter(name: str, cfg: omegaconf.DictConfig)[source]#
Add an Adapter module to this module.
- Parameters
name – A globally unique name for the adapter. Will be used to access, enable and disable adapters.
cfg – A DictConfig or Dataclass that contains at the bare minimum __target__ to instantiate a new Adapter module.
- batch_concat_states(batch_states: List[List[torch.Tensor]]) List[torch.Tensor] [source]#
Concatenate a batch of decoder state to a packed state.
- Parameters
batch_states (list) – batch of decoder states B x ([L x (H)], [L x (H)])
- Returns
- decoder states
(L x B x H, L x B x H)
- Return type
(tuple)
- batch_copy_states(old_states: List[torch.Tensor], new_states: List[torch.Tensor], ids: List[int], value: Optional[float] = None) List[torch.Tensor] [source]#
Copy states from new state to old state at certain indices.
- Parameters
old_states (list) – packed decoder states (L x B x H, L x B x H)
new_states – packed decoder states (L x B x H, L x B x H)
ids (list) – List of indices to copy states at.
value (optional float) – If a value should be copied instead of a state slice, a float should be provided
- Returns
- batch of decoder states with partial copy at ids (or a specific value).
(L x B x H, L x B x H)
- batch_initialize_states(batch_states: List[torch.Tensor], decoder_states: List[List[torch.Tensor]])[source]#
Create batch of decoder states.
- Parameters
batch_states (list) – batch of decoder states ([L x (B, H)], [L x (B, H)])
decoder_states (list of list) – list of decoder states [B x ([L x (1, H)], [L x (1, H)])]
- Returns
- batch of decoder states
([L x (B, H)], [L x (B, H)])
- Return type
batch_states (tuple)
- batch_score_hypothesis(hypotheses: List[Hypothesis], cache: Dict[Tuple[int], Any], batch_states: List[torch.Tensor]) Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor] [source]#
Used for batched beam search algorithms. Similar to score_hypothesis method.
- Parameters
hypothesis – List of Hypotheses. Refer to rnnt_utils.Hypothesis.
cache – Dict which contains a cache to avoid duplicate computations.
batch_states – List of torch.Tensor which represent the states of the RNN for this batch. Each state is of shape [L, B, H]
- Returns
b_y is a torch.Tensor of shape [B, 1, H] representing the scores of the last tokens in the Hypotheses. b_state is a list of list of RNN states, each of shape [L, B, H].
Represented as B x List[states].
lm_token is a list of the final integer tokens of the hypotheses in the batch.
- Return type
Returns a tuple (b_y, b_states, lm_tokens) such that
- batch_select_state(batch_states: List[torch.Tensor], idx: int) List[List[torch.Tensor]] [source]#
Get decoder state from batch of states, for given id.
- Parameters
batch_states (list) – batch of decoder states ([L x (B, H)], [L x (B, H)])
idx (int) – index to extract state from batch of states
- Returns
- decoder states for given id
([L x (1, H)], [L x (1, H)])
- Return type
(tuple)
- initialize_state(y: torch.Tensor) List[torch.Tensor] [source]#
Initialize the state of the RNN layers, with same dtype and device as input y.
- Parameters
y – A torch.Tensor whose device the generated states will be placed on.
- Returns
- List of torch.Tensor, each of shape [L, B, H], where
L = Number of RNN layers B = Batch size H = Hidden size of RNN.
- input_example(max_batch=1, max_dim=1)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- predict(y: Optional[torch.Tensor] = None, state: Optional[List[torch.Tensor]] = None, add_sos: bool = True, batch_size: Optional[int] = None) Tuple[torch.Tensor, List[torch.Tensor]] [source]#
Stateful prediction of scores and state for a (possibly null) tokenset. This method takes various cases into consideration : - No token, no state - used for priming the RNN - No token, state provided - used for blank token scoring - Given token, states - used for scores + new states
Here: B - batch size U - label length H - Hidden dimension size of RNN L - Number of RNN layers
- Parameters
y – Optional torch tensor of shape [B, U] of dtype long which will be passed to the Embedding. If None, creates a zero tensor of shape [B, 1, H] which mimics output of pad-token on EmbeddiNg.
state – An optional list of states for the RNN. Eg: For LSTM, it is the state list length is 2. Each state must be a tensor of shape [L, B, H]. If None, and during training mode and random_state_sampling is set, will sample a normal distribution tensor of the above shape. Otherwise, None will be passed to the RNN.
add_sos – bool flag, whether a zero vector describing a “start of signal” token should be prepended to the above “y” tensor. When set, output size is (B, U + 1, H).
batch_size – An optional int, specifying the batch size of the y tensor. Can be infered if y and state is None. But if both are None, then batch_size cannot be None.
- Returns
A tuple (g, hid) such that -
- If add_sos is False:
g: (B, U, H) hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
h (tensor), shape (L, B, H) c (tensor), shape (L, B, H)
- If add_sos is True:
g: (B, U + 1, H) hid: (h, c) where h is the final sequence hidden state and c is the final cell state:
h (tensor), shape (L, B, H) c (tensor), shape (L, B, H)
- score_hypothesis(hypothesis: Hypothesis, cache: Dict[Tuple[int], Any]) Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor] [source]#
Similar to the predict() method, instead this method scores a Hypothesis during beam search. Hypothesis is a dataclass representing one hypothesis in a Beam Search.
- Parameters
hypothesis – Refer to rnnt_utils.Hypothesis.
cache – Dict which contains a cache to avoid duplicate computations.
- Returns
y is a torch.Tensor of shape [1, 1, H] representing the score of the last token in the Hypothesis. state is a list of RNN states, each of shape [L, 1, H]. lm_token is the final integer token of the hypothesis.
- Return type
Returns a tuple (y, states, lm_token) such that
- class nemo.collections.asr.modules.StatelessTransducerDecoder(*args: Any, **kwargs: Any)[source]#
Bases:
AbstractRNNTDecoder
,Exportable
A Stateless Neural Network Transducer Decoder / Prediction Network. An RNN-T Decoder/Prediction stateless network that simply takes concatenation of embeddings of the history tokens as the output.
- Parameters
prednet –
A dict-like object which contains the following key-value pairs. pred_hidden: int specifying the hidden dimension of the prediction net.
dropout: float, set to 0.0 by default. Optional dropout applied at the end of the final LSTM RNN layer.
vocab_size – int, specifying the vocabulary size of the embedding layer of the Prediction network, excluding the RNNT blank token.
context_size – int, specifying the size of the history context used for this decoder.
normalization_mode – Can be either None, ‘layer’. By default, is set to None. Defines the type of normalization applied to the RNN layer.
- batch_concat_states(batch_states: List[List[torch.Tensor]]) List[torch.Tensor] [source]#
Concatenate a batch of decoder state to a packed state.
- Parameters
batch_states (list) – batch of decoder states B x ([(C)]
- Returns
- decoder states
[(B x C)]
- Return type
(tuple)
- batch_copy_states(old_states: List[torch.Tensor], new_states: List[torch.Tensor], ids: List[int], value: Optional[float] = None) List[torch.Tensor] [source]#
Copy states from new state to old state at certain indices.
- Parameters
old_states – packed decoder states single element list of (B x C)
new_states – packed decoder states single element list of (B x C)
ids (list) – List of indices to copy states at.
value (optional float) – If a value should be copied instead of a state slice, a float should be provided
- Returns
- batch of decoder states with partial copy at ids (or a specific value).
(B x C)
- batch_initialize_states(batch_states: List[torch.Tensor], decoder_states: List[List[torch.Tensor]])[source]#
Create batch of decoder states.
- Parameters
batch_states (list) – batch of decoder states ([(B, H)])
decoder_states (list of list) – list of decoder states [B x ([(1, C)]]
- Returns
- batch of decoder states
([(B, C)])
- Return type
batch_states (tuple)
- batch_score_hypothesis(hypotheses: List[Hypothesis], cache: Dict[Tuple[int], Any], batch_states: List[torch.Tensor]) Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor] [source]#
Used for batched beam search algorithms. Similar to score_hypothesis method.
- Parameters
hypothesis – List of Hypotheses. Refer to rnnt_utils.Hypothesis.
cache – Dict which contains a cache to avoid duplicate computations.
batch_states – List of torch.Tensor which represent the states of the RNN for this batch. Each state is of shape [L, B, H]
- Returns
b_y is a torch.Tensor of shape [B, 1, H] representing the scores of the last tokens in the Hypotheses. b_state is a list of list of RNN states, each of shape [L, B, H].
Represented as B x List[states].
lm_token is a list of the final integer tokens of the hypotheses in the batch.
- Return type
Returns a tuple (b_y, b_states, lm_tokens) such that
- batch_select_state(batch_states: List[torch.Tensor], idx: int) List[List[torch.Tensor]] [source]#
Get decoder state from batch of states, for given id.
- Parameters
batch_states (list) – batch of decoder states [(B, C)]
idx (int) – index to extract state from batch of states
- Returns
- decoder states for given id
[(C)]
- Return type
(tuple)
- initialize_state(y: torch.Tensor) List[torch.Tensor] [source]#
Initialize the state of the RNN layers, with same dtype and device as input y.
- Parameters
y – A torch.Tensor whose device the generated states will be placed on.
- Returns
- List of torch.Tensor, each of shape [L, B, H], where
L = Number of RNN layers B = Batch size H = Hidden size of RNN.
- input_example(max_batch=1, max_dim=1)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- predict(y: Optional[torch.Tensor] = None, state: Optional[torch.Tensor] = None, add_sos: bool = True, batch_size: Optional[int] = None) Tuple[torch.Tensor, List[torch.Tensor]] [source]#
Stateful prediction of scores and state for a tokenset.
Here: B - batch size U - label length C - context size for stateless decoder D - total embedding size
- Parameters
y – Optional torch tensor of shape [B, U] of dtype long which will be passed to the Embedding. If None, creates a zero tensor of shape [B, 1, D] which mimics output of pad-token on Embedding.
state – An optional one-element list of one tensor. The tensor is used to store previous context labels. The tensor uses type long and is of shape [B, C].
add_sos – bool flag, whether a zero vector describing a “start of signal” token should be prepended to the above “y” tensor. When set, output size is (B, U + 1, D).
batch_size – An optional int, specifying the batch size of the y tensor. Can be infered if y and state is None. But if both are None, then batch_size cannot be None.
- Returns
A tuple (g, state) such that -
- If add_sos is False:
g: (B, U, D) state: [(B, C)] storing the history context including the new words in y.
- If add_sos is True:
g: (B, U + 1, D) state: [(B, C)] storing the history context including the new words in y.
- score_hypothesis(hypothesis: Hypothesis, cache: Dict[Tuple[int], Any]) Tuple[torch.Tensor, List[torch.Tensor], torch.Tensor] [source]#
Similar to the predict() method, instead this method scores a Hypothesis during beam search. Hypothesis is a dataclass representing one hypothesis in a Beam Search.
- Parameters
hypothesis – Refer to rnnt_utils.Hypothesis.
cache – Dict which contains a cache to avoid duplicate computations.
- Returns
y is a torch.Tensor of shape [1, 1, H] representing the score of the last token in the Hypothesis. state is a list of RNN states, each of shape [L, 1, H]. lm_token is the final integer token of the hypothesis.
- Return type
Returns a tuple (y, states, lm_token) such that
- class nemo.collections.asr.modules.RNNTJoint(*args: Any, **kwargs: Any)[source]#
Bases:
AbstractRNNTJoint
,Exportable
,AdapterModuleMixin
A Recurrent Neural Network Transducer Joint Network (RNN-T Joint Network). An RNN-T Joint network, comprised of a feedforward model.
- Parameters
jointnet –
A dict-like object which contains the following key-value pairs. encoder_hidden: int specifying the hidden dimension of the encoder net. pred_hidden: int specifying the hidden dimension of the prediction net. joint_hidden: int specifying the hidden dimension of the joint net activation: Activation function used in the joint step. Can be one of
[‘relu’, ‘tanh’, ‘sigmoid’].
Optionally, it may also contain the following: dropout: float, set to 0.0 by default. Optional dropout applied at the end of the joint net.
num_classes – int, specifying the vocabulary size that the joint network must predict, excluding the RNNT blank token.
vocabulary – Optional list of strings/tokens that comprise the vocabulary of the joint network. Unused and kept only for easy access for character based encoding RNNT models.
log_softmax – Optional bool, set to None by default. If set as None, will compute the log_softmax() based on the value provided.
preserve_memory –
Optional bool, set to False by default. If the model crashes due to the memory intensive joint step, one might try this flag to empty the tensor cache in pytorch.
Warning: This will make the forward-backward pass much slower than normal. It also might not fix the OOM if the GPU simply does not have enough memory to compute the joint.
fuse_loss_wer –
Optional bool, set to False by default.
Fuses the joint forward, loss forward and wer forward steps. In doing so, it trades of speed for memory conservation by creating sub-batches of the provided batch of inputs, and performs Joint forward, loss forward and wer forward (optional), all on sub-batches, then collates results to be exactly equal to results from the entire batch.
When this flag is set, prior to calling forward, the fields loss and wer (either one) must be set using the RNNTJoint.set_loss() or RNNTJoint.set_wer() methods.
Further, when this flag is set, the following argument fused_batch_size must be provided as a non negative integer. This value refers to the size of the sub-batch.
When the flag is set, the input and output signature of forward() of this method changes. Input - in addition to encoder_outputs (mandatory argument), the following arguments can be provided.
decoder_outputs (optional). Required if loss computation is required.
encoder_lengths (required)
transcripts (optional). Required for wer calculation.
transcript_lengths (optional). Required for wer calculation.
compute_wer (bool, default false). Whether to compute WER or not for the fused batch.
- Output - instead of the usual joint log prob tensor, the following results can be returned.
loss (optional). Returned if decoder_outputs, transcripts and transript_lengths are not None.
- wer_numerator + wer_denominator (optional). Returned if transcripts, transcripts_lengths are provided
and compute_wer is set.
fused_batch_size – Optional int, required if fuse_loss_wer flag is set. Determines the size of the sub-batches. Should be any value below the actual batch size per GPU.
- add_adapter(name: str, cfg: omegaconf.DictConfig)[source]#
Add an Adapter module to this module.
- Parameters
name – A globally unique name for the adapter. Will be used to access, enable and disable adapters.
cfg – A DictConfig or Dataclass that contains at the bare minimum __target__ to instantiate a new Adapter module.
- property disabled_deployment_input_names#
Implement this method to return a set of input names disabled for export
- forward(encoder_outputs: torch.Tensor, decoder_outputs: Optional[torch.Tensor], encoder_lengths: Optional[torch.Tensor] = None, transcripts: Optional[torch.Tensor] = None, transcript_lengths: Optional[torch.Tensor] = None, compute_wer: bool = False) Union[torch.Tensor, List[Optional[torch.Tensor]]] [source]#
- property fuse_loss_wer#
- property fused_batch_size#
- input_example(max_batch=1, max_dim=8192)[source]#
Generates input examples for tracing etc. :returns: A tuple of input examples.
- property input_types#
Returns definitions of module input ports.
- joint(f: torch.Tensor, g: torch.Tensor) torch.Tensor [source]#
Compute the joint step of the network.
Here, B = Batch size T = Acoustic model timesteps U = Target sequence length H1, H2 = Hidden dimensions of the Encoder / Decoder respectively H = Hidden dimension of the Joint hidden step. V = Vocabulary size of the Decoder (excluding the RNNT blank token).
Note
The implementation of this model is slightly modified from the original paper. The original paper proposes the following steps : (enc, dec) -> Expand + Concat + Sum [B, T, U, H1+H2] -> Forward through joint hidden [B, T, U, H] – *1 *1 -> Forward through joint final [B, T, U, V + 1].
We instead split the joint hidden into joint_hidden_enc and joint_hidden_dec and act as follows: enc -> Forward through joint_hidden_enc -> Expand [B, T, 1, H] – *1 dec -> Forward through joint_hidden_dec -> Expand [B, 1, U, H] – *2 (*1, *2) -> Sum [B, T, U, H] -> Forward through joint final [B, T, U, V + 1].
- Parameters
f – Output of the Encoder model. A torch.Tensor of shape [B, T, H1]
g – Output of the Decoder model. A torch.Tensor of shape [B, U, H2]
- Returns
Logits / log softmaxed tensor of shape (B, T, U, V + 1).
- property loss#
- property num_classes_with_blank#
- property num_extra_outputs#
- property output_types#
Returns definitions of module output ports.
- property wer#
- class nemo.collections.asr.modules.SampledRNNTJoint(*args: Any, **kwargs: Any)[source]#
Bases:
RNNTJoint
A Sampled Recurrent Neural Network Transducer Joint Network (RNN-T Joint Network). An RNN-T Joint network, comprised of a feedforward model, where the vocab size will be sampled instead of computing the full vocabulary joint.
- Parameters
jointnet –
A dict-like object which contains the following key-value pairs. encoder_hidden: int specifying the hidden dimension of the encoder net. pred_hidden: int specifying the hidden dimension of the prediction net. joint_hidden: int specifying the hidden dimension of the joint net activation: Activation function used in the joint step. Can be one of
[‘relu’, ‘tanh’, ‘sigmoid’].
Optionally, it may also contain the following: dropout: float, set to 0.0 by default. Optional dropout applied at the end of the joint net.
num_classes – int, specifying the vocabulary size that the joint network must predict, excluding the RNNT blank token.
n_samples – int, specifies the number of tokens to sample from the vocabulary space, excluding the RNNT blank token. If a given value is larger than the entire vocabulary size, then the full vocabulary will be used.
vocabulary – Optional list of strings/tokens that comprise the vocabulary of the joint network. Unused and kept only for easy access for character based encoding RNNT models.
log_softmax – Optional bool, set to None by default. If set as None, will compute the log_softmax() based on the value provided.
preserve_memory –
Optional bool, set to False by default. If the model crashes due to the memory intensive joint step, one might try this flag to empty the tensor cache in pytorch.
Warning: This will make the forward-backward pass much slower than normal. It also might not fix the OOM if the GPU simply does not have enough memory to compute the joint.
fuse_loss_wer –
Optional bool, set to False by default.
Fuses the joint forward, loss forward and wer forward steps. In doing so, it trades of speed for memory conservation by creating sub-batches of the provided batch of inputs, and performs Joint forward, loss forward and wer forward (optional), all on sub-batches, then collates results to be exactly equal to results from the entire batch.
When this flag is set, prior to calling forward, the fields loss and wer (either one) must be set using the RNNTJoint.set_loss() or RNNTJoint.set_wer() methods.
Further, when this flag is set, the following argument fused_batch_size must be provided as a non negative integer. This value refers to the size of the sub-batch.
When the flag is set, the input and output signature of forward() of this method changes. Input - in addition to encoder_outputs (mandatory argument), the following arguments can be provided.
decoder_outputs (optional). Required if loss computation is required.
encoder_lengths (required)
transcripts (optional). Required for wer calculation.
transcript_lengths (optional). Required for wer calculation.
compute_wer (bool, default false). Whether to compute WER or not for the fused batch.
- Output - instead of the usual joint log prob tensor, the following results can be returned.
loss (optional). Returned if decoder_outputs, transcripts and transript_lengths are not None.
- wer_numerator + wer_denominator (optional). Returned if transcripts, transcripts_lengths are provided
and compute_wer is set.
fused_batch_size – Optional int, required if fuse_loss_wer flag is set. Determines the size of the sub-batches. Should be any value below the actual batch size per GPU.
- forward(encoder_outputs: torch.Tensor, decoder_outputs: Optional[torch.Tensor], encoder_lengths: Optional[torch.Tensor] = None, transcripts: Optional[torch.Tensor] = None, transcript_lengths: Optional[torch.Tensor] = None, compute_wer: bool = False) Union[torch.Tensor, List[Optional[torch.Tensor]]] [source]#
- sampled_joint(f: torch.Tensor, g: torch.Tensor, transcript: torch.Tensor, transcript_lengths: torch.Tensor) torch.Tensor [source]#
Compute the sampled joint step of the network.
# Reference - [Memory-Efficient Training of RNN-Transducer with Sampled Softmax](https://arxiv.org/abs/2203.16868)
Here, B = Batch size T = Acoustic model timesteps U = Target sequence length H1, H2 = Hidden dimensions of the Encoder / Decoder respectively H = Hidden dimension of the Joint hidden step. V = Vocabulary size of the Decoder (excluding the RNNT blank token). S = Sample size of vocabulary.
Note
The implementation of this joint model is slightly modified from the original paper. The original paper proposes the following steps : (enc, dec) -> Expand + Concat + Sum [B, T, U, H1+H2] -> Forward through joint hidden [B, T, U, H] – *1 *1 -> Forward through joint final [B, T, U, V + 1].
We instead split the joint hidden into joint_hidden_enc and joint_hidden_dec and act as follows: enc -> Forward through joint_hidden_enc -> Expand [B, T, 1, H] – *1 dec -> Forward through joint_hidden_dec -> Expand [B, 1, U, H] – *2 (*1, *2) -> Sum [B, T, U, H] -> Sample Vocab V_Pos (for target tokens) and V_Neg -> (V_Neg is sampled not uniformly by as a rand permutation of all vocab tokens, then eliminate all Intersection(V_Pos, V_Neg) common tokens to avoid duplication of loss) -> Concat new Vocab V_Sampled = Union(V_Pos, V_Neg) -> Forward partially through the joint final to create [B, T, U, V_Sampled]
- Parameters
f – Output of the Encoder model. A torch.Tensor of shape [B, T, H1]
g – Output of the Decoder model. A torch.Tensor of shape [B, U, H2]
transcript – Batch of transcripts. A torch.Tensor of shape [B, U]
transcript_lengths – Batch of lengths of the transcripts. A torch.Tensor of shape [B]
- Returns
Logits / log softmaxed tensor of shape (B, T, U, V + 1).
Parts#
- class nemo.collections.asr.parts.submodules.jasper.JasperBlock(*args: Any, **kwargs: Any)[source]#
Bases:
Module
,AdapterModuleMixin
,AccessMixin
Constructs a single “Jasper” block. With modified parameters, also constructs other blocks for models such as QuartzNet and Citrinet.
For Jasper : separable flag should be False
For QuartzNet : separable flag should be True
For Citrinet : separable flag and se flag should be True
Note that above are general distinctions, each model has intricate differences that expand over multiple such blocks.
For further information about the differences between models which use JasperBlock, please review the configs for ASR models found in the ASR examples directory.
- Parameters
inplanes – Number of input channels.
planes – Number of output channels.
repeat – Number of repeated sub-blocks (R) for this block.
kernel_size – Convolution kernel size across all repeated sub-blocks.
kernel_size_factor – Floating point scale value that is multiplied with kernel size, then rounded down to nearest odd integer to compose the kernel size. Defaults to 1.0.
stride – Stride of the convolutional layers.
dilation – Integer which defined dilation factor of kernel. Note that when dilation > 1, stride must be equal to 1.
padding – String representing type of padding. Currently only supports “same” padding, which symmetrically pads the input tensor with zeros.
dropout – Floating point value, determins percentage of output that is zeroed out.
activation – String representing activation functions. Valid activation functions are : {“hardtanh”: nn.Hardtanh, “relu”: nn.ReLU, “selu”: nn.SELU, “swish”: Swish}. Defaults to “relu”.
residual – Bool that determined whether a residual branch should be added or not. All residual branches are constructed using a pointwise convolution kernel, that may or may not perform strided convolution depending on the parameter residual_mode.
groups – Number of groups for Grouped Convolutions. Defaults to 1.
separable – Bool flag that describes whether Time-Channel depthwise separable convolution should be constructed, or ordinary convolution should be constructed.
heads – Number of “heads” for the masked convolution. Defaults to -1, which disables it.
normalization – String that represents type of normalization performed. Can be one of “batch”, “group”, “instance” or “layer” to compute BatchNorm1D, GroupNorm1D, InstanceNorm or LayerNorm (which are special cases of GroupNorm1D).
norm_groups – Number of groups used for GroupNorm (if normalization == “group”).
residual_mode – String argument which describes whether the residual branch should be simply added (“add”) or should first stride, then add (“stride_add”). Required when performing stride on parallel branch as well as utilizing residual add.
residual_panes – Number of residual panes, used for Jasper-DR models. Please refer to the paper.
conv_mask – Bool flag which determines whether to utilize masked convolutions or not. In general, it should be set to True.
se – Bool flag that determines whether Squeeze-and-Excitation layer should be used.
se_reduction_ratio – Integer value, which determines to what extend the hidden dimension of the SE intermediate step should be reduced. Larger values reduce number of parameters, but also limit the effectiveness of SE layers.
se_context_window – Integer value determining the number of timesteps that should be utilized in order to compute the averaged context window. Defaults to -1, which means it uses global context - such that all timesteps are averaged. If any positive integer is used, it will utilize limited context window of that size.
se_interpolation_mode – String used for interpolation mode of timestep dimension for SE blocks. Used only if context window is > 1. The modes available for resizing are: nearest, linear (3D-only), bilinear, area.
stride_last – Bool flag that determines whether all repeated blocks should stride at once, (stride of S^R when this flag is False) or just the last repeated block should stride (stride of S when this flag is True).
future_context –
Int value that determins how many “right” / “future” context frames will be utilized when calculating the output of the conv kernel. All calculations are done for odd kernel sizes only.
By default, this is -1, which is recomputed as the symmetric padding case.
When future_context >= 0, will compute the asymmetric padding as follows : (left context, right context) = [K - 1 - future_context, future_context]
Determining an exact formula to limit future context is dependent on global layout of the model. As such, we provide both “local” and “global” guidelines below.
Local context limit (should always be enforced) - future context should be <= half the kernel size for any given layer - future context > kernel size defaults to symmetric kernel - future context of layer = number of future frames * width of each frame (dependent on stride)
Global context limit (should be carefully considered) - future context should be layed out in an ever reducing pattern. Initial layers should restrict future context less than later layers, since shallow depth (and reduced stride) means each frame uses less amounts of future context. - Beyond a certain point, future context should remain static for a given stride level. This is the upper bound of the amount of future context that can be provided to the model on a global scale. - future context is calculated (roughly) as - (2 ^ stride) * (K // 2) number of future frames. This resultant value should be bound to some global maximum number of future seconds of audio (in ms).
Note: In the special case where K < future_context, it is assumed that the kernel is too small to limit its future context, so symmetric padding is used instead.
Note: There is no explicit limitation on the amount of future context used, as long as K > future_context constraint is maintained. This might lead to cases where future_context is more than half the actual kernel size K! In such cases, the conv layer is utilizing more of the future context than its current and past context to compute the output. While this is possible to do, it is not recommended and the layer will raise a warning to notify the user of such cases. It is advised to simply use symmetric padding for such cases.
Example: Say we have a model that performs 8x stride and receives spectrogram frames with stride of 0.01s. Say we wish to upper bound future context to 80 ms.
Layer ID, Kernel Size, Stride, Future Context, Global Context 0, K=5, S=1, FC=8, GC= 2 * (2^0) = 2 * 0.01 ms (special case, K < FC so use symmetric pad) 1, K=7, S=1, FC=3, GC= 3 * (2^0) = 3 * 0.01 ms (note that symmetric pad here uses 3 FC frames!) 2, K=11, S=2, FC=4, GC= 4 * (2^1) = 8 * 0.01 ms (note that symmetric pad here uses 5 FC frames!) 3, K=15, S=1, FC=4, GC= 4 * (2^1) = 8 * 0.01 ms (note that symmetric pad here uses 7 FC frames!) 4, K=21, S=2, FC=2, GC= 2 * (2^2) = 8 * 0.01 ms (note that symmetric pad here uses 10 FC frames!) 5, K=25, S=2, FC=1, GC= 1 * (2^3) = 8 * 0.01 ms (note that symmetric pad here uses 14 FC frames!) 6, K=29, S=1, FC=1, GC= 1 * (2^3) = 8 * 0.01 ms …
quantize – Bool flag whether to quantize the Convolutional blocks.
layer_idx (int, optional) – can be specified to allow layer output capture for InterCTC loss. Defaults to -1.
- forward(input_: Tuple[List[torch.Tensor], Optional[torch.Tensor]]) Tuple[List[torch.Tensor], Optional[torch.Tensor]] [source]#
Forward pass of the module.
- Parameters
input – The input is a tuple of two values - the preprocessed audio signal as well as the lengths of the audio signal. The audio signal is padded to the shape [B, D, T] and the lengths are a torch vector of length B.
- Returns
The output of the block after processing the input through repeat number of sub-blocks, as well as the lengths of the encoded audio after padding/striding.
Mixins#
- class nemo.collections.asr.parts.mixins.mixins.ASRBPEMixin[source]#
Bases:
ABC
ASR BPE Mixin class that sets up a Tokenizer via a config
This mixin class adds the method _setup_tokenizer(…), which can be used by ASR models which depend on subword tokenization.
- The setup_tokenizer method adds the following parameters to the class -
tokenizer_cfg: The resolved config supplied to the tokenizer (with dir and type arguments).
tokenizer_dir: The directory path to the tokenizer vocabulary + additional metadata.
tokenizer_type: The type of the tokenizer. Currently supports bpe and wpe, as well as agg.
vocab_path: Resolved path to the vocabulary text file.
In addition to these variables, the method will also instantiate and preserve a tokenizer (subclass of TokenizerSpec) if successful, and assign it to self.tokenizer.
The mixin also supports aggregate tokenizers, which consist of ordinary, monolingual tokenizers. If a conversion between a monolongual and an aggregate tokenizer (or vice versa) is detected, all registered artifacts will be cleaned up.
- AGGREGATE_TOKENIZERS_DICT_PREFIX = 'langs'#
- class nemo.collections.asr.parts.mixins.mixins.ASRModuleMixin[source]#
Bases:
ASRAdapterModelMixin
ASRModuleMixin is a mixin class added to ASR models in order to add methods that are specific to a particular instantiation of a module inside of an ASRModel.
Each method should first check that the module is present within the subclass, and support additional functionality if the corresponding module is present.
- change_attention_model(self_attention_model: Optional[str] = None, att_context_size: Optional[List[int]] = None, update_config: bool = True)[source]#
Update the self_attention_model if function is available in encoder.
- Parameters
self_attention_model (str) –
type of the attention layer and positional encoding ‘rel_pos’: relative positional embedding and Transformer-XL ‘rel_pos_local_attn’: relative positional embedding and Transformer-XL with local attention using
overlapping windows. Attention context is determined by att_context_size parameter.
’abs_pos’: absolute positional embedding and Transformer If None is provided, the self_attention_model isn’t changed. Defauts to None.
att_context_size (List[int]) – List of 2 ints corresponding to left and right attention context sizes, or None to keep as it is. Defauts to None.
update_config (bool) – Whether to update the config or not with the new attention model. Defaults to True.
- change_conv_asr_se_context_window(context_window: int, update_config: bool = True)[source]#
Update the context window of the SqueezeExcitation module if the provided model contains an encoder which is an instance of ConvASREncoder.
- Parameters
context_window –
An integer representing the number of input timeframes that will be used to compute the context. Each timeframe corresponds to a single window stride of the STFT features.
Say the window_stride = 0.01s, then a context window of 128 represents 128 * 0.01 s of context to compute the Squeeze step.
update_config – Whether to update the config or not with the new context window.
- change_subsampling_conv_chunking_factor(subsampling_conv_chunking_factor: int, update_config: bool = True)[source]#
Update the conv_chunking_factor (int) if function is available in encoder. Default is 1 (auto) Set it to -1 (disabled) or to a specific value (power of 2) if you OOM in the conv subsampling layers
- Parameters
conv_chunking_factor (int) –
- conformer_stream_step(processed_signal: torch.Tensor, processed_signal_length: Optional[torch.Tensor] = None, cache_last_channel: Optional[torch.Tensor] = None, cache_last_time: Optional[torch.Tensor] = None, cache_last_channel_len: Optional[torch.Tensor] = None, keep_all_outputs: bool = True, previous_hypotheses: Optional[List[Hypothesis]] = None, previous_pred_out: Optional[torch.Tensor] = None, drop_extra_pre_encoded: Optional[int] = None, return_transcription: bool = True, return_log_probs: bool = False)[source]#
It simulates a forward step with caching for streaming purposes. It supports the ASR models where their encoder supports streaming like Conformer. :param processed_signal: the input audio signals :param processed_signal_length: the length of the audios :param cache_last_channel: the cache tensor for last channel layers like MHA :param cache_last_channel_len: engths for cache_last_channel :param cache_last_time: the cache tensor for last time layers like convolutions :param keep_all_outputs: if set to True, would not drop the extra outputs specified by encoder.streaming_cfg.valid_out_len :param previous_hypotheses: the hypotheses from the previous step for RNNT models :param previous_pred_out: the predicted outputs from the previous step for CTC models :param drop_extra_pre_encoded: number of steps to drop from the beginning of the outputs after the downsampling module. This can be used if extra paddings are added on the left side of the input. :param return_transcription: whether to decode and return the transcriptions. It can not get disabled for Transducer models. :param return_log_probs: whether to return the log probs, only valid for ctc model
- Returns
the greedy predictions from the decoder all_hyp_or_transcribed_texts: the decoder hypotheses for Transducer models and the transcriptions for CTC models cache_last_channel_next: the updated tensor cache for last channel layers to be used for next streaming step cache_last_time_next: the updated tensor cache for last time layers to be used for next streaming step cache_last_channel_next_len: the updated lengths for cache_last_channel best_hyp: the best hypotheses for the Transducer models
log_probs: the logits tensor of current streaming chunk, only returned when return_log_probs=True encoded_len: the length of the output log_probs + history chunk log_probs, only returned when return_log_probs=True
- Return type
greedy_predictions
- transcribe_simulate_cache_aware_streaming(paths2audio_files: List[str], batch_size: int = 4, logprobs: bool = False, return_hypotheses: bool = False, online_normalization: bool = False)#
- Parameters
paths2audio_files – (a list) of paths to audio files.
batch_size – (int) batch size to use during inference. Bigger will result in better throughput performance but would use more memory.
logprobs – (bool) pass True to get log probabilities instead of transcripts.
return_hypotheses – (bool) Either return hypotheses or text With hypotheses can do some postprocessing like getting timestamp or rescoring
online_normalization – (bool) Perform normalization on the run per chunk.
- Returns
A list of transcriptions (or raw log probabilities if logprobs is True) in the same order as paths2audio_files
- class nemo.collections.asr.parts.mixins.interctc_mixin.InterCTCMixin[source]#
Bases:
object
Adds utilities for computing interCTC loss from https://arxiv.org/abs/2102.03216.
To use, make sure encoder accesses
interctc['capture_layers']
property in the AccessMixin and registersinterctc/layer_output_X
andinterctc/layer_length_X
for all layers that we want to get loss from. Additionally, specify the following config parameters to set up loss:interctc: # can use different values loss_weights: [0.3] apply_at_layers: [8]
Then call
self.setup_interctc(ctc_decoder_name, ctc_loss_name, ctc_wer_name)
in the init methodself.add_interctc_losses
after computing regular loss.self.finalize_interctc_metrics(metrics, outputs, prefix="val_")
in the multi_validation_epoch_end method.self.finalize_interctc_metrics(metrics, outputs, prefix="test_")
in the multi_test_epoch_end method.
- add_interctc_losses(loss_value: torch.Tensor, transcript: torch.Tensor, transcript_len: torch.Tensor, compute_wer: bool, compute_loss: bool = True, log_wer_num_denom: bool = False, log_prefix: str = '') Tuple[Optional[torch.Tensor], Dict] [source]#
Adding interCTC losses if required.
Will also register loss/wer metrics in the returned dictionary.
- Parameters
loss_value (torch.Tensor) – regular loss tensor (will add interCTC loss to it).
transcript (torch.Tensor) – current utterance transcript.
transcript_len (torch.Tensor) – current utterance transcript length.
compute_wer (bool) – whether to compute WER for the current utterance. Should typically be True for validation/test and only True for training if current batch WER should be logged.
compute_loss (bool) – whether to compute loss for the current utterance. Should always be True in training and almost always True in validation, unless all other losses are disabled as well. Defaults to True.
log_wer_num_denom (bool) – if True, will additionally log WER num/denom in the returned metrics dictionary. Should always be True for validation/test to allow correct metrics aggregation. Should always be False for training. Defaults to False.
log_prefix (str) – prefix added to all log values. Should be
""
for training and"val_"
for validation. Defaults to “”.
- Returns
tuple of new loss tensor and dictionary with logged metrics.
- Return type
tuple[Optional[torch.Tensor], Dict]
- finalize_interctc_metrics(metrics: Dict, outputs: List[Dict], prefix: str)[source]#
Finalizes InterCTC WER and loss metrics for logging purposes.
Should be called inside
multi_validation_epoch_end
(withprefix="val_"
) ormulti_test_epoch_end
(withprefix="test_"
).Note that
metrics
dictionary is going to be updated in-place.
- get_captured_interctc_tensors() List[Tuple[torch.Tensor, torch.Tensor]] [source]#
Returns a list of captured tensors from encoder: tuples of (output, length).
Will additionally apply
ctc_decoder
to the outputs.
- get_interctc_param(param_name)[source]#
Either directly get parameter from
self._interctc_params
or call getattr with the corresponding name.
- set_interctc_param(param_name, param_value)[source]#
Setting the parameter to the
self._interctc_params
dictionary.Raises an error if trying to set decoder, loss or wer as those should always come from the main class.
- setup_interctc(decoder_name, loss_name, wer_name)[source]#
Sets up all interctc-specific parameters and checks config consistency.
Caller has to specify names of attributes to perform CTC-specific WER, decoder and loss computation. They will be looked up in the class state with
getattr
.The reason we get the names and look up object later is because those objects might change without re-calling the setup of this class. So we always want to look up the most up-to-date object instead of “caching” it here.
Datasets#
Character Encoding Datasets#
- class nemo.collections.asr.data.audio_to_text.AudioToCharDataset(*args: Any, **kwargs: Any)[source]#
Bases:
_AudioTextDataset
Dataset that loads tensors via a json file containing paths to audio files, transcripts, and durations (in seconds). Each new line is a different sample. Example below: {“audio_filepath”: “/path/to/audio.wav”, “text_filepath”: “/path/to/audio.txt”, “duration”: 23.147} … {“audio_filepath”: “/path/to/audio.wav”, “text”: “the transcription”, “offset”: 301.75, “duration”: 0.82, “utt”: “utterance_id”, “ctm_utt”: “en_4156”, “side”: “A”}
- Parameters
manifest_filepath – Path to manifest json as described above. Can be comma-separated paths.
labels – String containing all the possible characters to map to
sample_rate (int) – Sample rate to resample loaded audio to
int_values (bool) – If true, load samples as 32-bit integers. Defauts to False.
augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor) – An AudioAugmentor object used to augment loaded audio
max_duration – If audio exceeds this length, do not include in dataset
min_duration – If audio is less than this length, do not include in dataset
max_utts – Limit number of utterances
blank_index – blank character index, default = -1
unk_index – unk_character index, default = -1
normalize – whether to normalize transcript text (default): True
bos_id – Id of beginning of sequence symbol to append if not None
eos_id – Id of end of sequence symbol to append if not None
return_sample_id (bool) – whether to return the sample_id as a part of each sample
channel_selector (int | Iterable[int] | str) – select a single channel or a subset of channels from multi-channel audio. If set to ‘average’, it performs averaging across channels. Disabled if set to None. Defaults to None. Uses zero-based indexing.
- property output_types: Optional[Dict[str, NeuralType]]#
Returns definitions of module output ports.
- class nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset(*args: Any, **kwargs: Any)[source]#
Bases:
_TarredAudioToTextDataset
A similar Dataset to the AudioToCharDataset, but which loads tarred audio files.
Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToCharDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball.
Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. ‘path/to/audio.tar’ or ‘path/to/audio_{1..100}.tar.gz’, or (2) a list of file paths that will not be brace-expanded, e.g. [‘audio_1.tar’, ‘audio_2.tar’, …].
See the WebDataset documentation for more information about accepted data and input formats.
If using multiple workers the number of shards should be divisible by world_size to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven!
Notice that a few arguments are different from the AudioToCharDataset; for example, shuffle (bool) has been replaced by shuffle_n (int).
Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line.
- Parameters
audio_tar_filepaths – Either a list of audio tarball filepaths, or a string (can be brace-expandable).
manifest_filepath (str) – Path to the manifest.
labels (list) – List of characters that can be output by the ASR model. For Jasper, this is the 28 character set {a-z ‘}. The CTC blank symbol is automatically added later for models using ctc.
sample_rate (int) – Sample rate to resample loaded audio to
int_values (bool) – If true, load samples as 32-bit integers. Defauts to False.
augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor) – An AudioAugmentor object used to augment loaded audio
shuffle_n (int) – How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0.
min_duration (float) – Dataset parameter. All training files which have a duration less than min_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to 0.1.
max_duration (float) – Dataset parameter. All training files which have a duration more than max_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to None.
blank_index (int) – Blank character index, defaults to -1.
unk_index (int) – Unknown character index, defaults to -1.
normalize (bool) – Dataset parameter. Whether to use automatic text cleaning. It is highly recommended to manually clean text for best results. Defaults to True.
trim (bool) – Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim(). Defaults to False.
bos_id (id) – Dataset parameter. Beginning of string symbol id used for seq2seq models. Defaults to None.
eos_id (id) – Dataset parameter. End of string symbol id used for seq2seq models. Defaults to None.
pad_id (id) – Token used to pad when collating samples in batches. If this is None, pads using 0s. Defaults to None.
shard_strategy (str) –
Tarred dataset shard distribution strategy chosen as a str value during ddp.
scatter: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime.
replicate: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of shuffle_n.
Warning
Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with
world_size
), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets.
global_rank (int) – Worker rank, used for partitioning shards. Defaults to 0.
world_size (int) – Total number of processes, used for partitioning shards. Defaults to 0.
return_sample_id (bool) – whether to return the sample_id as a part of each sample
Text-to-Text Datasets for Hybrid ASR-TTS models#
- class nemo.collections.asr.data.text_to_text.TextToTextDataset(*args: Any, **kwargs: Any)[source]#
Bases:
TextToTextDatasetBase
,Dataset
Text-to-Text Map-style Dataset for hybrid ASR-TTS models
Subword Encoding Datasets#
- class nemo.collections.asr.data.audio_to_text.AudioToBPEDataset(*args: Any, **kwargs: Any)[source]#
Bases:
_AudioTextDataset
Dataset that loads tensors via a json file containing paths to audio files, transcripts, and durations (in seconds). Each new line is a different sample. Example below: {“audio_filepath”: “/path/to/audio.wav”, “text_filepath”: “/path/to/audio.txt”, “duration”: 23.147} … {“audio_filepath”: “/path/to/audio.wav”, “text”: “the transcription”, “offset”: 301.75, “duration”: 0.82, “utt”: “utterance_id”, “ctm_utt”: “en_4156”, “side”: “A”}
In practice, the dataset and manifest used for character encoding and byte pair encoding are exactly the same. The only difference lies in how the dataset tokenizes the text in the manifest.
- Parameters
manifest_filepath – Path to manifest json as described above. Can be comma-separated paths.
tokenizer – A subclass of the Tokenizer wrapper found in the common collection, nemo.collections.common.tokenizers.TokenizerSpec. ASR Models support a subset of all available tokenizers.
sample_rate (int) – Sample rate to resample loaded audio to
int_values (bool) – If true, load samples as 32-bit integers. Defauts to False.
augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor) – An AudioAugmentor object used to augment loaded audio
max_duration – If audio exceeds this length, do not include in dataset
min_duration – If audio is less than this length, do not include in dataset
max_utts – Limit number of utterances
trim – Whether to trim silence segments
use_start_end_token – Boolean which dictates whether to add [BOS] and [EOS] tokens to beginning and ending of speech respectively.
return_sample_id (bool) – whether to return the sample_id as a part of each sample
channel_selector (int | Iterable[int] | str) – select a single channel or a subset of channels from multi-channel audio. If set to ‘average’, it performs averaging across channels. Disabled if set to None. Defaults to None. Uses zero-based indexing.
- property output_types: Optional[Dict[str, NeuralType]]#
Returns definitions of module output ports.
- class nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset(*args: Any, **kwargs: Any)[source]#
Bases:
_TarredAudioToTextDataset
A similar Dataset to the AudioToBPEDataset, but which loads tarred audio files.
Accepts a single comma-separated JSON manifest file (in the same style as for the AudioToBPEDataset), as well as the path(s) to the tarball(s) containing the wav files. Each line of the manifest should contain the information for one audio file, including at least the transcript and name of the audio file within the tarball.
Valid formats for the audio_tar_filepaths argument include: (1) a single string that can be brace-expanded, e.g. ‘path/to/audio.tar’ or ‘path/to/audio_{1..100}.tar.gz’, or (2) a list of file paths that will not be brace-expanded, e.g. [‘audio_1.tar’, ‘audio_2.tar’, …].
See the WebDataset documentation for more information about accepted data and input formats.
If using multiple workers the number of shards should be divisible by world_size to ensure an even split among workers. If it is not divisible, logging will give a warning but training will proceed. In addition, if using mutiprocessing, each shard MUST HAVE THE SAME NUMBER OF ENTRIES after filtering is applied. We currently do not check for this, but your program may hang if the shards are uneven!
Notice that a few arguments are different from the AudioToBPEDataset; for example, shuffle (bool) has been replaced by shuffle_n (int).
Additionally, please note that the len() of this DataLayer is assumed to be the length of the manifest after filtering. An incorrect manifest length may lead to some DataLoader issues down the line.
- Parameters
audio_tar_filepaths – Either a list of audio tarball filepaths, or a string (can be brace-expandable).
manifest_filepath (str) – Path to the manifest.
tokenizer (TokenizerSpec) – Either a Word Piece Encoding tokenizer (BERT), or a Sentence Piece Encoding tokenizer (BPE). The CTC blank symbol is automatically added later for models using ctc.
sample_rate (int) – Sample rate to resample loaded audio to
int_values (bool) – If true, load samples as 32-bit integers. Defauts to False.
augmentor (nemo.collections.asr.parts.perturb.AudioAugmentor) – An AudioAugmentor object used to augment loaded audio
shuffle_n (int) – How many samples to look ahead and load to be shuffled. See WebDataset documentation for more details. Defaults to 0.
min_duration (float) – Dataset parameter. All training files which have a duration less than min_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to 0.1.
max_duration (float) – Dataset parameter. All training files which have a duration more than max_duration are dropped. Note: Duration is read from the manifest JSON. Defaults to None.
trim (bool) – Whether to use trim silence from beginning and end of audio signal using librosa.effects.trim(). Defaults to False.
use_start_end_token – Boolean which dictates whether to add [BOS] and [EOS] tokens to beginning and ending of speech respectively.
pad_id (id) – Token used to pad when collating samples in batches. If this is None, pads using 0s. Defaults to None.
shard_strategy (str) –
Tarred dataset shard distribution strategy chosen as a str value during ddp.
scatter: The default shard strategy applied by WebDataset, where each node gets a unique set of shards, which are permanently pre-allocated and never changed at runtime.
replicate: Optional shard strategy, where each node gets all of the set of shards available in the tarred dataset, which are permanently pre-allocated and never changed at runtime. The benefit of replication is that it allows each node to sample data points from the entire dataset independently of other nodes, and reduces dependence on value of shuffle_n.
Warning
Replicated strategy allows every node to sample the entire set of available tarfiles, and therefore more than one node may sample the same tarfile, and even sample the same data points! As such, there is no assured guarantee that all samples in the dataset will be sampled at least once during 1 epoch. Scattered strategy, on the other hand, on specific occasions (when the number of shards is not divisible with
world_size
), will not sample the entire dataset. For these reasons it is not advisable to use tarred datasets as validation or test datasets.
global_rank (int) – Worker rank, used for partitioning shards. Defaults to 0.
world_size (int) – Total number of processes, used for partitioning shards. Defaults to 0.
return_sample_id (bool) – whether to return the sample_id as a part of each sample
Audio Preprocessors#
- class nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor(*args: Any, **kwargs: Any)[source]#
Bases:
AudioPreprocessor
,Exportable
Featurizer module that converts wavs to mel spectrograms.
- Parameters
sample_rate (int) – Sample rate of the input audio data. Defaults to 16000
window_size (float) – Size of window for fft in seconds Defaults to 0.02
window_stride (float) – Stride of window for fft in seconds Defaults to 0.01
n_window_size (int) – Size of window for fft in samples Defaults to None. Use one of window_size or n_window_size.
n_window_stride (int) – Stride of window for fft in samples Defaults to None. Use one of window_stride or n_window_stride.
window (str) – Windowing function for fft. can be one of [‘hann’, ‘hamming’, ‘blackman’, ‘bartlett’] Defaults to “hann”
normalize (str) – Can be one of [‘per_feature’, ‘all_features’]; all other options disable feature normalization. ‘all_features’ normalizes the entire spectrogram to be mean 0 with std 1. ‘pre_features’ normalizes per channel / freq instead. Defaults to “per_feature”
n_fft (int) – Length of FT window. If None, it uses the smallest power of 2 that is larger than n_window_size. Defaults to None
preemph (float) – Amount of pre emphasis to add to audio. Can be disabled by passing None. Defaults to 0.97
features (int) – Number of mel spectrogram freq bins to output. Defaults to 64
lowfreq (int) – Lower bound on mel basis in Hz. Defaults to 0
highfreq (int) – Lower bound on mel basis in Hz. Defaults to None
log (bool) – Log features. Defaults to True
log_zero_guard_type (str) – Need to avoid taking the log of zero. There are two options: “add” or “clamp”. Defaults to “add”.
log_zero_guard_value (float, or str) – Add or clamp requires the number to add with or clamp to. log_zero_guard_value can either be a float or “tiny” or “eps”. torch.finfo is used if “tiny” or “eps” is passed. Defaults to 2**-24.
dither (float) – Amount of white-noise dithering. Defaults to 1e-5
pad_to (int) – Ensures that the output size of the time dimension is a multiple of pad_to. Defaults to 16
frame_splicing (int) – Defaults to 1
exact_pad (bool) – If True, sets stft center to False and adds padding, such that num_frames = audio_length // hop_length. Defaults to False.
pad_value (float) – The value that shorter mels are padded with. Defaults to 0
mag_power (float) – The power that the linear spectrogram is raised to prior to multiplication with mel basis. Defaults to 2 for a power spec
rng – Random number generator
nb_augmentation_prob (float) – Probability with which narrowband augmentation would be applied to samples in the batch. Defaults to 0.0
nb_max_freq (int) – Frequency above which all frequencies will be masked for narrowband augmentation. Defaults to 4000
use_torchaudio – Whether to use the torchaudio implementation.
mel_norm – Normalization used for mel filterbank weights. Defaults to ‘slaney’ (area normalization)
stft_exact_pad – Deprecated argument, kept for compatibility with older checkpoints.
stft_conv – Deprecated argument, kept for compatibility with older checkpoints.
- property filter_banks#
- input_example(max_batch: int = 8, max_dim: int = 32000, min_length: int = 200)[source]#
Override this method if random inputs won’t work :returns: A tuple sample of valid input data.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- processed_signal:
0: AxisType(BatchTag) 1: AxisType(MelSpectrogramSignalTag) 2: AxisType(ProcessedTimeTag)
- processed_length:
0: AxisType(BatchTag)
- classmethod restore_from(restore_path: str)[source]#
Restores model instance (weights and configuration) from a .nemo file
- Parameters
restore_path – path to .nemo file from which model should be instantiated
override_config_path – path to a yaml config that will override the internal config file or an OmegaConf / DictConfig object representing the model config.
map_location – Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise.
strict – Passed to load_state_dict. By default True
return_config – If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model.
trainer – An optional Trainer object, passed to the model constructor.
save_restore_connector – An optional SaveRestoreConnector object that defines the implementation of the restore_from() method.
- save_to(save_path: str)[source]#
Standardized method to save a tarfile containing the checkpoint, config, and any additional artifacts. Implemented via
nemo.core.connectors.save_restore_connector.SaveRestoreConnector.save_to()
.- Parameters
save_path – str, path to where the file should be saved.
- class nemo.collections.asr.modules.AudioToMFCCPreprocessor(*args: Any, **kwargs: Any)[source]#
Bases:
AudioPreprocessor
Preprocessor that converts wavs to MFCCs. Uses torchaudio.transforms.MFCC.
- Parameters
sample_rate – The sample rate of the audio. Defaults to 16000.
window_size – Size of window for fft in seconds. Used to calculate the win_length arg for mel spectrogram. Defaults to 0.02
window_stride – Stride of window for fft in seconds. Used to caculate the hop_length arg for mel spect. Defaults to 0.01
n_window_size – Size of window for fft in samples Defaults to None. Use one of window_size or n_window_size.
n_window_stride – Stride of window for fft in samples Defaults to None. Use one of window_stride or n_window_stride.
window – Windowing function for fft. can be one of [‘hann’, ‘hamming’, ‘blackman’, ‘bartlett’, ‘none’, ‘null’]. Defaults to ‘hann’
n_fft – Length of FT window. If None, it uses the smallest power of 2 that is larger than n_window_size. Defaults to None
lowfreq (int) – Lower bound on mel basis in Hz. Defaults to 0
highfreq (int) – Lower bound on mel basis in Hz. Defaults to None
n_mels – Number of mel filterbanks. Defaults to 64
n_mfcc – Number of coefficients to retain Defaults to 64
dct_type – Type of discrete cosine transform to use
norm – Type of norm to use
log – Whether to use log-mel spectrograms instead of db-scaled. Defaults to True.
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- classmethod restore_from(restore_path: str)[source]#
Restores model instance (weights and configuration) from a .nemo file
- Parameters
restore_path – path to .nemo file from which model should be instantiated
override_config_path – path to a yaml config that will override the internal config file or an OmegaConf / DictConfig object representing the model config.
map_location – Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise.
strict – Passed to load_state_dict. By default True
return_config – If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model.
trainer – An optional Trainer object, passed to the model constructor.
save_restore_connector – An optional SaveRestoreConnector object that defines the implementation of the restore_from() method.
- save_to(save_path: str)[source]#
Standardized method to save a tarfile containing the checkpoint, config, and any additional artifacts. Implemented via
nemo.core.connectors.save_restore_connector.SaveRestoreConnector.save_to()
.- Parameters
save_path – str, path to where the file should be saved.
Audio Augmentors#
- class nemo.collections.asr.modules.SpectrogramAugmentation(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
Performs time and freq cuts in one of two ways. SpecAugment zeroes out vertical and horizontal sections as described in SpecAugment (https://arxiv.org/abs/1904.08779). Arguments for use with SpecAugment are freq_masks, time_masks, freq_width, and time_width. SpecCutout zeroes out rectangulars as described in Cutout (https://arxiv.org/abs/1708.04552). Arguments for use with Cutout are rect_masks, rect_freq, and rect_time.
- Parameters
freq_masks (int) – how many frequency segments should be cut. Defaults to 0.
time_masks (int) – how many time segments should be cut Defaults to 0.
freq_width (int) – maximum number of frequencies to be cut in one segment. Defaults to 10.
time_width (int) – maximum number of time steps to be cut in one segment Defaults to 10.
rect_masks (int) – how many rectangular masks should be cut Defaults to 0.
rect_freq (int) – maximum size of cut rectangles along the frequency dimension Defaults to 5.
rect_time (int) – maximum size of cut rectangles along the time dimension Defaults to 25.
- property input_types#
Returns definitions of module input types
- property output_types#
Returns definitions of module output types
- class nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation(*args: Any, **kwargs: Any)[source]#
Bases:
NeuralModule
Pad or Crop the incoming Spectrogram to a certain shape.
- Parameters
audio_length (int) – the final number of timesteps that is required. The signal will be either padded or cropped temporally to this size.
- forward#
- property input_types#
Returns definitions of module output ports.
- property output_types#
Returns definitions of module output ports.
- classmethod restore_from(restore_path: str)[source]#
Restores model instance (weights and configuration) from a .nemo file
- Parameters
restore_path – path to .nemo file from which model should be instantiated
override_config_path – path to a yaml config that will override the internal config file or an OmegaConf / DictConfig object representing the model config.
map_location – Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise.
strict – Passed to load_state_dict. By default True
return_config – If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model.
trainer – An optional Trainer object, passed to the model constructor.
save_restore_connector – An optional SaveRestoreConnector object that defines the implementation of the restore_from() method.
- save_to(save_path: str)[source]#
Standardized method to save a tarfile containing the checkpoint, config, and any additional artifacts. Implemented via
nemo.core.connectors.save_restore_connector.SaveRestoreConnector.save_to()
.- Parameters
save_path – str, path to where the file should be saved.
- class nemo.collections.asr.parts.preprocessing.perturb.SpeedPerturbation(sr, resample_type, min_speed_rate=0.9, max_speed_rate=1.1, num_rates=5, rng=None)[source]#
Bases:
Perturbation
Performs Speed Augmentation by re-sampling the data to a different sampling rate, which does not preserve pitch.
Note: This is a very slow operation for online augmentation. If space allows, it is preferable to pre-compute and save the files to augment the dataset.
- Parameters
sr – Original sampling rate.
resample_type – Type of resampling operation that will be performed. For better speed using resampy’s fast resampling method, use resample_type=’kaiser_fast’. For high-quality resampling, set resample_type=’kaiser_best’. To use scipy.signal.resample, set resample_type=’fft’ or resample_type=’scipy’
min_speed_rate – Minimum sampling rate modifier.
max_speed_rate – Maximum sampling rate modifier.
num_rates – Number of discrete rates to allow. Can be a positive or negative integer. If a positive integer greater than 0 is provided, the range of speed rates will be discretized into num_rates values. If a negative integer or 0 is provided, the full range of speed rates will be sampled uniformly. Note: If a positive integer is provided and the resultant discretized range of rates contains the value ‘1.0’, then those samples with rate=1.0, will not be augmented at all and simply skipped. This is to unnecessary augmentation and increase computation time. Effective augmentation chance in such a case is = prob * (num_rates - 1 / num_rates) * 100`% chance where `prob is the global probability of a sample being augmented.
rng – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.TimeStretchPerturbation(min_speed_rate=0.9, max_speed_rate=1.1, num_rates=5, n_fft=512, rng=None)[source]#
Bases:
Perturbation
Time-stretch an audio series by a fixed rate while preserving pitch, based on [1, 2].
Note: This is a simplified implementation, intended primarily for reference and pedagogical purposes. It makes no attempt to handle transients, and is likely to produce audible artifacts.
Reference [1] [Ellis, D. P. W. “A phase vocoder in Matlab.” Columbia University, 2002.] (http://www.ee.columbia.edu/~dpwe/resources/matlab/pvoc/) [2] [librosa.effects.time_stretch] (https://librosa.github.io/librosa/generated/librosa.effects.time_stretch.html)
- Parameters
min_speed_rate – Minimum sampling rate modifier.
max_speed_rate – Maximum sampling rate modifier.
num_rates – Number of discrete rates to allow. Can be a positive or negative integer. If a positive integer greater than 0 is provided, the range of speed rates will be discretized into num_rates values. If a negative integer or 0 is provided, the full range of speed rates will be sampled uniformly. Note: If a positive integer is provided and the resultant discretized range of rates contains the value ‘1.0’, then those samples with rate=1.0, will not be augmented at all and simply skipped. This is to avoid unnecessary augmentation and increase computation time. Effective augmentation chance in such a case is = prob * (num_rates - 1 / num_rates) * 100`% chance where `prob is the global probability of a sample being augmented.
n_fft – Number of fft filters to be computed.
rng – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.GainPerturbation(min_gain_dbfs=-10, max_gain_dbfs=10, rng=None)[source]#
Bases:
Perturbation
Applies random gain to the audio.
- Parameters
min_gain_dbfs (float) – Min gain level in dB
max_gain_dbfs (float) – Max gain level in dB
rng (int) – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.ImpulsePerturbation(manifest_path=None, audio_tar_filepaths=None, shuffle_n=128, normalize_impulse=False, shift_impulse=False, rng=None)[source]#
Bases:
Perturbation
Convolves audio with a Room Impulse Response.
- Parameters
manifest_path (list) – Manifest file for RIRs
audio_tar_filepaths (list) – Tar files, if RIR audio files are tarred
shuffle_n (int) – Shuffle parameter for shuffling buffered files from the tar files
normalize_impulse (bool) – Normalize impulse response to zero mean and amplitude 1
shift_impulse (bool) – Shift impulse response to adjust for delay at the beginning
rng (int) – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.ShiftPerturbation(min_shift_ms=-5.0, max_shift_ms=5.0, rng=None)[source]#
Bases:
Perturbation
Perturbs audio by shifting the audio in time by a random amount between min_shift_ms and max_shift_ms. The final length of the audio is kept unaltered by padding the audio with zeros.
- Parameters
min_shift_ms (float) – Minimum time in milliseconds by which audio will be shifted
max_shift_ms (float) – Maximum time in milliseconds by which audio will be shifted
rng (int) – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.NoisePerturbation(manifest_path=None, min_snr_db=10, max_snr_db=50, max_gain_db=300.0, rng=None, audio_tar_filepaths=None, shuffle_n=100, orig_sr=16000)[source]#
Bases:
Perturbation
Perturbation that adds noise to input audio.
- Parameters
manifest_path (str) – Manifest file with paths to noise files
min_snr_db (float) – Minimum SNR of audio after noise is added
max_snr_db (float) – Maximum SNR of audio after noise is added
max_gain_db (float) – Maximum gain that can be applied on the noise sample
audio_tar_filepaths (list) – Tar files, if noise audio files are tarred
shuffle_n (int) – Shuffle parameter for shuffling buffered files from the tar files
orig_sr (int) – Original sampling rate of the noise files
rng (int) – Random seed. Default is None
- property orig_sr#
- perturb(data, ref_mic=0)[source]#
- Parameters
data (AudioSegment) – audio data
ref_mic (int) – reference mic index for scaling multi-channel audios
- perturb_with_foreground_noise(data, noise, data_rms=None, max_noise_dur=2, max_additions=1, ref_mic=0)[source]#
- Parameters
data (AudioSegment) – audio data
noise (AudioSegment) – noise data
data_rms (Union[float, List[float]) – rms_db for data input
max_noise_dur – (float): max noise duration
max_additions (int) – number of times for adding noise
ref_mic (int) – reference mic index for scaling multi-channel audios
- class nemo.collections.asr.parts.preprocessing.perturb.WhiteNoisePerturbation(min_level=-90, max_level=-46, rng=None)[source]#
Bases:
Perturbation
Perturbation that adds white noise to an audio file in the training dataset.
- Parameters
min_level (int) – Minimum level in dB at which white noise should be added
max_level (int) – Maximum level in dB at which white noise should be added
rng (int) – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.RirAndNoisePerturbation(rir_manifest_path=None, rir_prob=0.5, noise_manifest_paths=None, noise_prob=1.0, min_snr_db=0, max_snr_db=50, rir_tar_filepaths=None, rir_shuffle_n=100, noise_tar_filepaths=None, apply_noise_rir=False, orig_sample_rate=None, max_additions=5, max_duration=2.0, bg_noise_manifest_paths=None, bg_noise_prob=1.0, bg_min_snr_db=10, bg_max_snr_db=50, bg_noise_tar_filepaths=None, bg_orig_sample_rate=None, rng=None)[source]#
Bases:
Perturbation
RIR augmentation with additive foreground and background noise. In this implementation audio data is augmented by first convolving the audio with a Room Impulse Response and then adding foreground noise and background noise at various SNRs. RIR, foreground and background noises should either be supplied with a manifest file or as tarred audio files (faster).
Different sets of noise audio files based on the original sampling rate of the noise. This is useful while training a mixed sample rate model. For example, when training a mixed model with 8 kHz and 16 kHz audio with a target sampling rate of 16 kHz, one would want to augment 8 kHz data with 8 kHz noise rather than 16 kHz noise.
- Parameters
rir_manifest_path – Manifest file for RIRs
rir_tar_filepaths – Tar files, if RIR audio files are tarred
rir_prob – Probability of applying a RIR
noise_manifest_paths – Foreground noise manifest path
min_snr_db – Min SNR for foreground noise
max_snr_db – Max SNR for background noise,
noise_tar_filepaths – Tar files, if noise files are tarred
apply_noise_rir – Whether to convolve foreground noise with a a random RIR
orig_sample_rate – Original sampling rate of foreground noise audio
max_additions – Max number of times foreground noise is added to an utterance,
max_duration – Max duration of foreground noise
bg_noise_manifest_paths – Background noise manifest path
bg_min_snr_db – Min SNR for background noise
bg_max_snr_db – Max SNR for background noise
bg_noise_tar_filepaths – Tar files, if noise files are tarred
bg_orig_sample_rate – Original sampling rate of background noise audio
rng – Random seed. Default is None
- class nemo.collections.asr.parts.preprocessing.perturb.TranscodePerturbation(codecs=None, rng=None)[source]#
Bases:
Perturbation
Audio codec augmentation. This implementation uses sox to transcode audio with low rate audio codecs, so users need to make sure that the installed sox version supports the codecs used here (G711 and amr-nb).
- Parameters
codecs (List[str]) – A list of codecs to be trancoded to. Default is None.
rng (int) – Random seed. Default is None.
Miscellaneous Classes#
CTC Decoding#
- class nemo.collections.asr.metrics.wer.CTCDecoding(decoding_cfg, vocabulary)[source]#
Bases:
AbstractCTCDecoding
Used for performing CTC auto-regressive / non-auto-regressive decoding of the logprobs for character based models.
- Parameters
decoding_cfg –
A dict-like object which contains the following key-value pairs. strategy: str value which represents the type of decoding that can occur.
Possible values are : - greedy (for greedy decoding). - beam (for DeepSpeed KenLM based decoding).
- compute_timestamps: A bool flag, which determines whether to compute the character/subword, or
word based timestamp mapping the output log-probabilities to discrite intervals of timestamps. The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
- ctc_timestamp_type: A str value, which represents the types of timestamps that should be calculated.
Can take the following values - “char” for character/subword time stamps, “word” for word level time stamps and “all” (default), for both character level and word level time stamps.
word_seperator: Str token representing the seperator between words.
- preserve_alignments: Bool flag which preserves the history of logprobs generated during
decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for logprobs in it. Here, logprobs is a torch.Tensors.
- confidence_cfg: A dict-like object which contains the following key-value pairs related to confidence
scores. In order to obtain hypotheses with confidence scores, please utilize ctc_decoder_predictions_tensor function with the preserve_frame_confidence flag set to True.
- preserve_frame_confidence: Bool flag which preserves the history of per-frame confidence scores
generated during decoding. When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, frame_confidence is a List of floats.
- preserve_token_confidence: Bool flag which preserves the history of per-token confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for token_confidence in it. Here, token_confidence is a List of floats.
The length of the list corresponds to the number of recognized tokens.
- preserve_word_confidence: Bool flag which preserves the history of per-word confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for word_confidence in it. Here, word_confidence is a List of floats.
The length of the list corresponds to the number of recognized words.
- exclude_blank: Bool flag indicating that blank token confidence scores are to be excluded
from the token_confidence.
- aggregation: Which aggregation type to use for collapsing per-token confidence into per-word confidence.
Valid options are mean, min, max, prod.
- measure_cfg: A dict-like object which contains the measure name and settings to compute per-frame
confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str).
Used if confidence_measure_cfg.name is set to entropy. Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
- batch_dim_index: Index of the batch dimension of
targets
andpredictions
parameters of ctc_decoder_predictions_tensor
methods. Can be either 0 or 1.
The config may further contain the following sub-dictionaries: “greedy”:
preserve_alignments: Same as above, overrides above value. compute_timestamps: Same as above, overrides above value. preserve_frame_confidence: Same as above, overrides above value. confidence_measure_cfg: Same as above, overrides confidence_cfg.measure_cfg.
- ”beam”:
- beam_size: int, defining the beam size for beam search. Must be >= 1.
If beam_size == 1, will perform cached greedy search. This might be slightly different results compared to the greedy search above.
- return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
hypotheses after beam search has concluded. This flag is set by default.
- beam_alpha: float, the strength of the Language model on the final score of a token.
final_score = acoustic_score + beam_alpha * lm_score + beam_beta * seq_length.
- beam_beta: float, the strength of the sequence length penalty on the final score of a token.
final_score = acoustic_score + beam_alpha * lm_score + beam_beta * seq_length.
- kenlm_path: str, path to a KenLM ARPA or .binary file (depending on the strategy chosen).
If the path is invalid (file is not found at path), will raise a deferred error at the moment of calculation of beam search, so that users may update / change the decoding strategy to point to the correct file.
blank_id – The id of the RNNT blank token.
- class nemo.collections.asr.metrics.wer_bpe.CTCBPEDecoding(decoding_cfg, tokenizer: TokenizerSpec)[source]#
Bases:
AbstractCTCDecoding
Used for performing CTC auto-regressive / non-auto-regressive decoding of the logprobs for subword based models.
- Parameters
decoding_cfg –
A dict-like object which contains the following key-value pairs. strategy: str value which represents the type of decoding that can occur.
Possible values are : - greedy (for greedy decoding). - beam (for DeepSpeed KenLM based decoding).
- compute_timestamps: A bool flag, which determines whether to compute the character/subword, or
word based timestamp mapping the output log-probabilities to discrite intervals of timestamps. The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
- ctc_timestamp_type: A str value, which represents the types of timestamps that should be calculated.
Can take the following values - “char” for character/subword time stamps, “word” for word level time stamps and “all” (default), for both character level and word level time stamps.
word_seperator: Str token representing the seperator between words.
- preserve_alignments: Bool flag which preserves the history of logprobs generated during
decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for logprobs in it. Here, logprobs is a torch.Tensors.
- confidence_cfg: A dict-like object which contains the following key-value pairs related to confidence
scores. In order to obtain hypotheses with confidence scores, please utilize ctc_decoder_predictions_tensor function with the preserve_frame_confidence flag set to True.
- preserve_frame_confidence: Bool flag which preserves the history of per-frame confidence scores
generated during decoding. When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, frame_confidence is a List of floats.
- preserve_token_confidence: Bool flag which preserves the history of per-token confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for token_confidence in it. Here, token_confidence is a List of floats.
The length of the list corresponds to the number of recognized tokens.
- preserve_word_confidence: Bool flag which preserves the history of per-word confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for word_confidence in it. Here, word_confidence is a List of floats.
The length of the list corresponds to the number of recognized words.
- exclude_blank: Bool flag indicating that blank token confidence scores are to be excluded
from the token_confidence.
- aggregation: Which aggregation type to use for collapsing per-token confidence into per-word confidence.
Valid options are mean, min, max, prod.
- measure_cfg: A dict-like object which contains the measure name and settings to compute per-frame
confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str).
Used if confidence_measure_cfg.name is set to entropy. Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
- batch_dim_index: Index of the batch dimension of
targets
andpredictions
parameters of ctc_decoder_predictions_tensor
methods. Can be either 0 or 1.
The config may further contain the following sub-dictionaries: “greedy”:
preserve_alignments: Same as above, overrides above value. compute_timestamps: Same as above, overrides above value. preserve_frame_confidence: Same as above, overrides above value. confidence_measure_cfg: Same as above, overrides confidence_cfg.measure_cfg.
- ”beam”:
- beam_size: int, defining the beam size for beam search. Must be >= 1.
If beam_size == 1, will perform cached greedy search. This might be slightly different results compared to the greedy search above.
- return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
hypotheses after beam search has concluded. This flag is set by default.
- beam_alpha: float, the strength of the Language model on the final score of a token.
final_score = acoustic_score + beam_alpha * lm_score + beam_beta * seq_length.
- beam_beta: float, the strength of the sequence length penalty on the final score of a token.
final_score = acoustic_score + beam_alpha * lm_score + beam_beta * seq_length.
- kenlm_path: str, path to a KenLM ARPA or .binary file (depending on the strategy chosen).
If the path is invalid (file is not found at path), will raise a deferred error at the moment of calculation of beam search, so that users may update / change the decoding strategy to point to the correct file.
tokenizer – NeMo tokenizer object, which inherits from TokenizerSpec.
- class nemo.collections.asr.parts.submodules.ctc_greedy_decoding.GreedyCTCInfer(blank_id: int, preserve_alignments: bool = False, compute_timestamps: bool = False, preserve_frame_confidence: bool = False, confidence_measure_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Bases:
Typing
,ConfidenceMeasureMixin
A greedy CTC decoder.
Provides a common abstraction for sample level and batch level greedy decoding.
- Parameters
blank_index – int index of the blank token. Can be 0 or len(vocabulary).
preserve_alignments – Bool flag which preserves the history of logprobs generated during decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for logprobs in it. Here, logprobs is a torch.Tensors.
compute_timestamps – A bool flag, which determines whether to compute the character/subword, or word based timestamp mapping the output log-probabilities to discrite intervals of timestamps. The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
preserve_frame_confidence – Bool flag which preserves the history of per-frame confidence scores generated during decoding. When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, frame_confidence is a List of floats.
confidence_measure_cfg –
A dict-like object which contains the measure name and settings to compute per-frame confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str). Used if confidence_measure_cfg.name is set to entropy.
- Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
- forward(decoder_output: torch.Tensor, decoder_lengths: torch.Tensor)[source]#
Returns a list of hypotheses given an input batch of the encoder hidden embedding. Output token is generated auto-repressively.
- Parameters
decoder_output – A tensor of size (batch, timesteps, features) or (batch, timesteps) (each timestep is a label).
decoder_lengths – list of int representing the length of each sequence output sequence.
- Returns
packed list containing batch number of sentences (Hypotheses).
- property input_types#
Returns definitions of module input ports.
- property output_types#
Returns definitions of module output ports.
- class nemo.collections.asr.parts.submodules.ctc_beam_decoding.BeamCTCInfer(blank_id: int, beam_size: int, search_type: str = 'default', return_best_hypothesis: bool = True, preserve_alignments: bool = False, compute_timestamps: bool = False, beam_alpha: float = 1.0, beam_beta: float = 0.0, kenlm_path: Optional[str] = None, flashlight_cfg: Optional[FlashlightConfig] = None, pyctcdecode_cfg: Optional[PyCTCDecodeConfig] = None)[source]#
Bases:
AbstractBeamCTCInfer
A greedy CTC decoder.
Provides a common abstraction for sample level and batch level greedy decoding.
- Parameters
blank_index – int index of the blank token. Can be 0 or len(vocabulary).
preserve_alignments – Bool flag which preserves the history of logprobs generated during decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for logprobs in it. Here, logprobs is a torch.Tensors.
compute_timestamps – A bool flag, which determines whether to compute the character/subword, or word based timestamp mapping the output log-probabilities to discrite intervals of timestamps. The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
- default_beam_search(x: torch.Tensor, out_len: torch.Tensor) List[Union[Hypothesis, NBestHypotheses]] #
Open Seq2Seq Beam Search Algorithm (DeepSpeed)
- Parameters
x – Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length, and V is the vocabulary size. The tensor contains log-probabilities.
out_len – Tensor of shape [B], contains lengths of each sequence in the batch.
- Returns
A list of NBestHypotheses objects, one for each sequence in the batch.
- flashlight_beam_search(x: torch.Tensor, out_len: torch.Tensor) List[Union[Hypothesis, NBestHypotheses]] #
Flashlight Beam Search Algorithm. Should support Char and Subword models.
- Parameters
x – Tensor of shape [B, T, V+1], where B is the batch size, T is the maximum sequence length, and V is the vocabulary size. The tensor contains log-probabilities.
out_len – Tensor of shape [B], contains lengths of each sequence in the batch.
- Returns
A list of NBestHypotheses objects, one for each sequence in the batch.
- forward(decoder_output: torch.Tensor, decoder_lengths: torch.Tensor) Tuple[List[Union[Hypothesis, NBestHypotheses]]] [source]#
Returns a list of hypotheses given an input batch of the encoder hidden embedding. Output token is generated auto-repressively.
- Parameters
decoder_output – A tensor of size (batch, timesteps, features).
decoder_lengths – list of int representing the length of each sequence output sequence.
- Returns
packed list containing batch number of sentences (Hypotheses).
RNNT Decoding#
- class nemo.collections.asr.metrics.rnnt_wer.RNNTDecoding(decoding_cfg, decoder, joint, vocabulary)[source]#
Bases:
AbstractRNNTDecoding
Used for performing RNN-T auto-regressive decoding of the Decoder+Joint network given the encoder state.
- Parameters
decoding_cfg –
A dict-like object which contains the following key-value pairs. strategy: str value which represents the type of decoding that can occur.
Possible values are : - greedy, greedy_batch (for greedy decoding). - beam, tsd, alsd (for beam search decoding).
- compute_hypothesis_token_set: A bool flag, which determines whether to compute a list of decoded
tokens as well as the decoded string. Default is False in order to avoid double decoding unless required.
- preserve_alignments: Bool flag which preserves the history of logprobs generated during
decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for logprobs in it. Here, alignments is a List of List of Tuple(Tensor (of length V + 1), Tensor(scalar, label after argmax)).
In order to obtain this hypothesis, please utilize rnnt_decoder_predictions_tensor function with the return_hypotheses flag set to True.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary. U is the number of target tokens for the current timestep Ti.
- confidence_cfg: A dict-like object which contains the following key-value pairs related to confidence
scores. In order to obtain hypotheses with confidence scores, please utilize rnnt_decoder_predictions_tensor function with the preserve_frame_confidence flag set to True.
- preserve_frame_confidence: Bool flag which preserves the history of per-frame confidence scores
generated during decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, alignments is a List of List of floats.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more confidence scores. U is the number of target tokens for the current timestep Ti.
- preserve_token_confidence: Bool flag which preserves the history of per-token confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for token_confidence in it. Here, token_confidence is a List of floats.
The length of the list corresponds to the number of recognized tokens.
- preserve_word_confidence: Bool flag which preserves the history of per-word confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for word_confidence in it. Here, word_confidence is a List of floats.
The length of the list corresponds to the number of recognized words.
- exclude_blank: Bool flag indicating that blank token confidence scores are to be excluded
from the token_confidence.
- aggregation: Which aggregation type to use for collapsing per-token confidence into per-word confidence.
Valid options are mean, min, max, prod.
- measure_cfg: A dict-like object which contains the measure name and settings to compute per-frame
confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str).
Used if confidence_measure_cfg.name is set to entropy. Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
The config may further contain the following sub-dictionaries: “greedy”:
- max_symbols: int, describing the maximum number of target tokens to decode per
timestep during greedy decoding. Setting to larger values allows longer sentences to be decoded, at the cost of increased execution time.
preserve_frame_confidence: Same as above, overrides above value.
confidence_measure_cfg: Same as above, overrides confidence_cfg.measure_cfg.
- ”beam”:
- beam_size: int, defining the beam size for beam search. Must be >= 1.
If beam_size == 1, will perform cached greedy search. This might be slightly different results compared to the greedy search above.
- score_norm: optional bool, whether to normalize the returned beam score in the hypotheses.
Set to True by default.
- return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
hypotheses after beam search has concluded. This flag is set by default.
- tsd_max_sym_exp: optional int, determines number of symmetric expansions of the target symbols
per timestep of the acoustic model. Larger values will allow longer sentences to be decoded, at increased cost to execution time.
- alsd_max_target_len: optional int or float, determines the potential maximum target sequence length.
If an integer is provided, it can decode sequences of that particular maximum length. If a float is provided, it can decode sequences of int(alsd_max_target_len * seq_len), where seq_len is the length of the acoustic model output (T).
- NOTE:
If a float is provided, it can be greater than 1! By default, a float of 2.0 is used so that a target sequence can be at most twice as long as the acoustic model output length T.
- maes_num_steps: Number of adaptive steps to take. From the paper, 2 steps is generally sufficient,
and can be reduced to 1 to improve decoding speed while sacrificing some accuracy. int > 0.
- maes_prefix_alpha: Maximum prefix length in prefix search. Must be an integer, and is advised to keep this as 1
in order to reduce expensive beam search cost later. int >= 0.
- maes_expansion_beta: Maximum number of prefix expansions allowed, in addition to the beam size.
Effectively, the number of hypothesis = beam_size + maes_expansion_beta. Must be an int >= 0, and affects the speed of inference since large values will perform large beam search in the next step.
- maes_expansion_gamma: Float pruning threshold used in the prune-by-value step when computing the expansions.
The default (2.3) is selected from the paper. It performs a comparison (max_log_prob - gamma <= log_prob[v]) where v is all vocabulary indices in the Vocab set and max_log_prob is the “most” likely token to be predicted. Gamma therefore provides a margin of additional tokens which can be potential candidates for expansion apart from the “most likely” candidate. Lower values will reduce the number of expansions (by increasing pruning-by-value, thereby improving speed but hurting accuracy). Higher values will increase the number of expansions (by reducing pruning-by-value, thereby reducing speed but potentially improving accuracy). This is a hyper parameter to be experimentally tuned on a validation set.
softmax_temperature: Scales the logits of the joint prior to computing log_softmax.
decoder – The Decoder/Prediction network module.
joint – The Joint network module.
vocabulary – The vocabulary (excluding the RNNT blank token) which will be used for decoding.
- decode_ids_to_langs(tokens: List[int]) List[str] [source]#
Decode a token id list into language ID (LID) list.
- Parameters
tokens – List of int representing the token ids.
- Returns
A list of decoded LIDS.
- decode_ids_to_tokens(tokens: List[int]) List[str] [source]#
Implemented by subclass in order to decode a token id list into a token list. A token list is the string representation of each token id.
- Parameters
tokens – List of int representing the token ids.
- Returns
A list of decoded tokens.
- class nemo.collections.asr.metrics.rnnt_wer_bpe.RNNTBPEDecoding(decoding_cfg, decoder, joint, tokenizer: TokenizerSpec)[source]#
Bases:
AbstractRNNTDecoding
Used for performing RNN-T auto-regressive decoding of the Decoder+Joint network given the encoder state.
- Parameters
decoding_cfg –
A dict-like object which contains the following key-value pairs. strategy: str value which represents the type of decoding that can occur.
Possible values are : - greedy, greedy_batch (for greedy decoding). - beam, tsd, alsd (for beam search decoding).
- compute_hypothesis_token_set: A bool flag, which determines whether to compute a list of decoded
tokens as well as the decoded string. Default is False in order to avoid double decoding unless required.
- preserve_alignments: Bool flag which preserves the history of logprobs generated during
decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for alignments in it. Here, alignments is a List of List of Tuple(Tensor (of length V + 1), Tensor(scalar, label after argmax)).
In order to obtain this hypothesis, please utilize rnnt_decoder_predictions_tensor function with the return_hypotheses flag set to True.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary. U is the number of target tokens for the current timestep Ti.
- compute_timestamps: A bool flag, which determines whether to compute the character/subword, or
word based timestamp mapping the output log-probabilities to discrete intervals of timestamps. The timestamps will be available in the returned Hypothesis.timestep as a dictionary.
- compute_langs: a bool flag, which allows to compute language id (LID) information per token,
word, and the entire sample (most likely language id). The LIDS will be available in the returned Hypothesis object as a dictionary
- rnnt_timestamp_type: A str value, which represents the types of timestamps that should be calculated.
Can take the following values - “char” for character/subword time stamps, “word” for word level time stamps and “all” (default), for both character level and word level time stamps.
word_seperator: Str token representing the seperator between words.
- preserve_frame_confidence: Bool flag which preserves the history of per-frame confidence scores
generated during decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, alignments is a List of List of ints.
- confidence_cfg: A dict-like object which contains the following key-value pairs related to confidence
scores. In order to obtain hypotheses with confidence scores, please utilize rnnt_decoder_predictions_tensor function with the preserve_frame_confidence flag set to True.
- preserve_frame_confidence: Bool flag which preserves the history of per-frame confidence scores
generated during decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, alignments is a List of List of floats.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more confidence scores. U is the number of target tokens for the current timestep Ti.
- preserve_token_confidence: Bool flag which preserves the history of per-token confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for token_confidence in it. Here, token_confidence is a List of floats.
The length of the list corresponds to the number of recognized tokens.
- preserve_word_confidence: Bool flag which preserves the history of per-word confidence scores
generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for word_confidence in it. Here, word_confidence is a List of floats.
The length of the list corresponds to the number of recognized words.
- exclude_blank: Bool flag indicating that blank token confidence scores are to be excluded
from the token_confidence.
- aggregation: Which aggregation type to use for collapsing per-token confidence into per-word confidence.
Valid options are mean, min, max, prod.
- measure_cfg: A dict-like object which contains the measure name and settings to compute per-frame
confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str).
Used if confidence_measure_cfg.name is set to entropy. Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
The config may further contain the following sub-dictionaries: “greedy”:
- max_symbols: int, describing the maximum number of target tokens to decode per
timestep during greedy decoding. Setting to larger values allows longer sentences to be decoded, at the cost of increased execution time.
preserve_frame_confidence: Same as above, overrides above value.
confidence_measure_cfg: Same as above, overrides confidence_cfg.measure_cfg.
- ”beam”:
- beam_size: int, defining the beam size for beam search. Must be >= 1.
If beam_size == 1, will perform cached greedy search. This might be slightly different results compared to the greedy search above.
- score_norm: optional bool, whether to normalize the returned beam score in the hypotheses.
Set to True by default.
- return_best_hypothesis: optional bool, whether to return just the best hypothesis or all of the
hypotheses after beam search has concluded.
- tsd_max_sym_exp: optional int, determines number of symmetric expansions of the target symbols
per timestep of the acoustic model. Larger values will allow longer sentences to be decoded, at increased cost to execution time.
- alsd_max_target_len: optional int or float, determines the potential maximum target sequence length.
If an integer is provided, it can decode sequences of that particular maximum length. If a float is provided, it can decode sequences of int(alsd_max_target_len * seq_len), where seq_len is the length of the acoustic model output (T).
- NOTE:
If a float is provided, it can be greater than 1! By default, a float of 2.0 is used so that a target sequence can be at most twice as long as the acoustic model output length T.
- maes_num_steps: Number of adaptive steps to take. From the paper, 2 steps is generally sufficient,
and can be reduced to 1 to improve decoding speed while sacrificing some accuracy. int > 0.
- maes_prefix_alpha: Maximum prefix length in prefix search. Must be an integer, and is advised to keep this as 1
in order to reduce expensive beam search cost later. int >= 0.
- maes_expansion_beta: Maximum number of prefix expansions allowed, in addition to the beam size.
Effectively, the number of hypothesis = beam_size + maes_expansion_beta. Must be an int >= 0, and affects the speed of inference since large values will perform large beam search in the next step.
- maes_expansion_gamma: Float pruning threshold used in the prune-by-value step when computing the expansions.
The default (2.3) is selected from the paper. It performs a comparison (max_log_prob - gamma <= log_prob[v]) where v is all vocabulary indices in the Vocab set and max_log_prob is the “most” likely token to be predicted. Gamma therefore provides a margin of additional tokens which can be potential candidates for expansion apart from the “most likely” candidate. Lower values will reduce the number of expansions (by increasing pruning-by-value, thereby improving speed but hurting accuracy). Higher values will increase the number of expansions (by reducing pruning-by-value, thereby reducing speed but potentially improving accuracy). This is a hyper parameter to be experimentally tuned on a validation set.
softmax_temperature: Scales the logits of the joint prior to computing log_softmax.
decoder – The Decoder/Prediction network module.
joint – The Joint network module.
tokenizer – The tokenizer which will be used for decoding.
- decode_hypothesis(hypotheses_list: List[Hypothesis]) List[Union[Hypothesis, NBestHypotheses]] [source]#
Decode a list of hypotheses into a list of strings. Overrides the super() method optionally adding lang information
- Parameters
hypotheses_list – List of Hypothesis.
- Returns
A list of strings.
- decode_ids_to_langs(tokens: List[int]) List[str] [source]#
Decode a token id list into language ID (LID) list.
- Parameters
tokens – List of int representing the token ids.
- Returns
A list of decoded LIDS.
- decode_ids_to_tokens(tokens: List[int]) List[str] [source]#
Implemented by subclass in order to decode a token id list into a token list. A token list is the string representation of each token id.
- Parameters
tokens – List of int representing the token ids.
- Returns
A list of decoded tokens.
- class nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyRNNTInfer(decoder_model: AbstractRNNTDecoder, joint_model: AbstractRNNTJoint, blank_index: int, max_symbols_per_step: Optional[int] = None, preserve_alignments: bool = False, preserve_frame_confidence: bool = False, confidence_measure_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Bases:
_GreedyRNNTInfer
A greedy transducer decoder.
Sequence level greedy decoding, performed auto-regressively.
- Parameters
decoder_model – rnnt_utils.AbstractRNNTDecoder implementation.
joint_model – rnnt_utils.AbstractRNNTJoint implementation.
blank_index – int index of the blank token. Can be 0 or len(vocabulary).
max_symbols_per_step – Optional int. The maximum number of symbols that can be added to a sequence in a single time step; if set to None then there is no limit.
preserve_alignments –
Bool flag which preserves the history of alignments generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for alignments in it. Here, alignments is a List of List of Tuple(Tensor (of length V + 1), Tensor(scalar, label after argmax)).
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary. U is the number of target tokens for the current timestep Ti.
preserve_frame_confidence –
Bool flag which preserves the history of per-frame confidence scores generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, frame_confidence is a List of List of floats.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more confidence scores. U is the number of target tokens for the current timestep Ti.
confidence_measure_cfg –
A dict-like object which contains the measure name and settings to compute per-frame confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str). Used if confidence_measure_cfg.name is set to entropy.
- Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
- forward(encoder_output: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[List[Hypothesis]] = None)[source]#
Returns a list of hypotheses given an input batch of the encoder hidden embedding. Output token is generated auto-regressively.
- Parameters
encoder_output – A tensor of size (batch, features, timesteps).
encoded_lengths – list of int representing the length of each sequence output sequence.
- Returns
packed list containing batch number of sentences (Hypotheses).
- class nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyBatchedRNNTInfer(decoder_model: AbstractRNNTDecoder, joint_model: AbstractRNNTJoint, blank_index: int, max_symbols_per_step: Optional[int] = None, preserve_alignments: bool = False, preserve_frame_confidence: bool = False, confidence_measure_cfg: Optional[omegaconf.DictConfig] = None)[source]#
Bases:
_GreedyRNNTInfer
A batch level greedy transducer decoder.
Batch level greedy decoding, performed auto-regressively.
- Parameters
decoder_model – rnnt_utils.AbstractRNNTDecoder implementation.
joint_model – rnnt_utils.AbstractRNNTJoint implementation.
blank_index – int index of the blank token. Can be 0 or len(vocabulary).
max_symbols_per_step – Optional int. The maximum number of symbols that can be added to a sequence in a single time step; if set to None then there is no limit.
preserve_alignments –
Bool flag which preserves the history of alignments generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for alignments in it. Here, alignments is a List of List of Tuple(Tensor (of length V + 1), Tensor(scalar, label after argmax)).
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary. U is the number of target tokens for the current timestep Ti.
preserve_frame_confidence –
Bool flag which preserves the history of per-frame confidence scores generated during greedy decoding (sample / batched). When set to true, the Hypothesis will contain the non-null value for frame_confidence in it. Here, frame_confidence is a List of List of floats.
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more confidence scores. U is the number of target tokens for the current timestep Ti.
confidence_measure_cfg –
A dict-like object which contains the measure name and settings to compute per-frame confidence scores.
- name: The measure name (str).
- Supported values:
’max_prob’ for using the maximum token probability as a confidence.
’entropy’ for using a normalized entropy of a log-likelihood vector.
- entropy_type: Which type of entropy to use (str). Used if confidence_measure_cfg.name is set to entropy.
- Supported values:
- ’gibbs’ for the (standard) Gibbs entropy. If the alpha (α) is provided,
the formula is the following: H_α = -sum_i((p^α_i)*log(p^α_i)). Note that for this entropy, the alpha should comply the following inequality: (log(V)+2-sqrt(log^2(V)+4))/(2*log(V)) <= α <= (1+log(V-1))/log(V-1) where V is the model vocabulary size.
- ’tsallis’ for the Tsallis entropy with the Boltzmann constant one.
Tsallis entropy formula is the following: H_α = 1/(α-1)*(1-sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/Tsallis_entropy
- ’renyi’ for the Rényi entropy.
Rényi entropy formula is the following: H_α = 1/(1-α)*log_2(sum_i(p^α_i)), where α is a parameter. When α == 1, it works like the Gibbs entropy. More: https://en.wikipedia.org/wiki/R%C3%A9nyi_entropy
- alpha: Power scale for logsoftmax (α for entropies). Here we restrict it to be > 0.
When the alpha equals one, scaling is not applied to ‘max_prob’, and any entropy type behaves like the Shannon entropy: H = -sum_i(p_i*log(p_i))
- entropy_norm: A mapping of the entropy value to the interval [0,1].
- Supported values:
’lin’ for using the linear mapping.
’exp’ for using exponential mapping with linear shift.
- forward(encoder_output: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[List[Hypothesis]] = None)[source]#
Returns a list of hypotheses given an input batch of the encoder hidden embedding. Output token is generated auto-regressively.
- Parameters
encoder_output – A tensor of size (batch, features, timesteps).
encoded_lengths – list of int representing the length of each sequence output sequence.
- Returns
packed list containing batch number of sentences (Hypotheses).
- class nemo.collections.asr.parts.submodules.rnnt_beam_decoding.BeamRNNTInfer(decoder_model: AbstractRNNTDecoder, joint_model: AbstractRNNTJoint, beam_size: int, search_type: str = 'default', score_norm: bool = True, return_best_hypothesis: bool = True, tsd_max_sym_exp_per_step: Optional[int] = 50, alsd_max_target_len: Union[int, float] = 1.0, nsc_max_timesteps_expansion: int = 1, nsc_prefix_alpha: int = 1, maes_num_steps: int = 2, maes_prefix_alpha: int = 1, maes_expansion_gamma: float = 2.3, maes_expansion_beta: int = 2, language_model: Optional[Dict[str, Any]] = None, softmax_temperature: float = 1.0, preserve_alignments: bool = False, ngram_lm_model: Optional[str] = None, ngram_lm_alpha: float = 0.0, hat_subtract_ilm: bool = False, hat_ilm_weight: float = 0.0)[source]#
Bases:
Typing
Beam Search implementation ported from ESPNet implementation - espnet/espnet
Sequence level beam decoding or batched-beam decoding, performed auto-repressively depending on the search type chosen.
- Parameters
decoder_model – rnnt_utils.AbstractRNNTDecoder implementation.
joint_model – rnnt_utils.AbstractRNNTJoint implementation.
beam_size –
number of beams for beam search. Must be a positive integer >= 1. If beam size is 1, defaults to stateful greedy search. This greedy search might result in slightly different results than the greedy results obtained by GreedyRNNTInfer due to implementation differences.
For accurate greedy results, please use GreedyRNNTInfer or GreedyBatchedRNNTInfer.
search_type (# The following arguments are specific to the chosen) –
str representing the type of beam search to perform. Must be one of [‘beam’, ‘tsd’, ‘alsd’]. ‘nsc’ is currently not supported.
Algoritm used: beam - basic beam search strategy. Larger beams generally result in better decoding,
however the time required for the search also grows steadily.
- tsd - time synchronous decoding. Please refer to the paper:
[Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040) for details on the algorithm implemented.
Time synchronous decoding (TSD) execution time grows by the factor T * max_symmetric_expansions. For longer sequences, T is greater, and can therefore take a long time for beams to obtain good results. This also requires greater memory to execute.
- alsd - alignment-length synchronous decoding. Please refer to the paper:
[Alignment-Length Synchronous Decoding for RNN Transducer](https://ieeexplore.ieee.org/document/9053040) for details on the algorithm implemented.
Alignment-length synchronous decoding (ALSD) execution time is faster than TSD, with growth factor of T + U_max, where U_max is the maximum target length expected during execution.
Generally, T + U_max < T * max_symmetric_expansions. However, ALSD beams are non-unique, therefore it is required to use larger beam sizes to achieve the same (or close to the same) decoding accuracy as TSD.
For a given decoding accuracy, it is possible to attain faster decoding via ALSD than TSD.
- maes = modified adaptive expansion searcn. Please refer to the paper:
[Accelerating RNN Transducer Inference via Adaptive Expansion Search](https://ieeexplore.ieee.org/document/9250505)
Modified Adaptive Synchronous Decoding (mAES) execution time is adaptive w.r.t the number of expansions (for tokens) required per timestep. The number of expansions can usually be constrained to 1 or 2, and in most cases 2 is sufficient.
This beam search technique can possibly obtain superior WER while sacrificing some evaluation time.
score_norm – bool, whether to normalize the scores of the log probabilities.
return_best_hypothesis – bool, decides whether to return a single hypothesis (the best out of N), or return all N hypothesis (sorted with best score first). The container class changes based this flag - When set to True (default), returns a single Hypothesis. When set to False, returns a NBestHypotheses container, which contains a list of Hypothesis.
search_type –
tsd_max_sym_exp_per_step – Used for search_type=tsd. The maximum symmetric expansions allowed per timestep during beam search. Larger values should be used to attempt decoding of longer sequences, but this in turn increases execution time and memory usage.
alsd_max_target_len – Used for search_type=alsd. The maximum expected target sequence length during beam search. Larger values allow decoding of longer sequences at the expense of execution time and memory.
stabilized. (# The following two flags are placeholders and unused until nsc implementation is) –
nsc_max_timesteps_expansion – Unused int.
nsc_prefix_alpha – Unused int.
flags (# mAES) –
maes_num_steps – Number of adaptive steps to take. From the paper, 2 steps is generally sufficient. int > 1.
maes_prefix_alpha – Maximum prefix length in prefix search. Must be an integer, and is advised to keep this as 1 in order to reduce expensive beam search cost later. int >= 0.
maes_expansion_beta – Maximum number of prefix expansions allowed, in addition to the beam size. Effectively, the number of hypothesis = beam_size + maes_expansion_beta. Must be an int >= 0, and affects the speed of inference since large values will perform large beam search in the next step.
maes_expansion_gamma – Float pruning threshold used in the prune-by-value step when computing the expansions. The default (2.3) is selected from the paper. It performs a comparison (max_log_prob - gamma <= log_prob[v]) where v is all vocabulary indices in the Vocab set and max_log_prob is the “most” likely token to be predicted. Gamma therefore provides a margin of additional tokens which can be potential candidates for expansion apart from the “most likely” candidate. Lower values will reduce the number of expansions (by increasing pruning-by-value, thereby improving speed but hurting accuracy). Higher values will increase the number of expansions (by reducing pruning-by-value, thereby reducing speed but potentially improving accuracy). This is a hyper parameter to be experimentally tuned on a validation set.
softmax_temperature – Scales the logits of the joint prior to computing log_softmax.
preserve_alignments –
Bool flag which preserves the history of alignments generated during beam decoding (sample). When set to true, the Hypothesis will contain the non-null value for alignments in it. Here, alignments is a List of List of Tensor (of length V + 1).
The length of the list corresponds to the Acoustic Length (T). Each value in the list (Ti) is a torch.Tensor (U), representing 1 or more targets from a vocabulary. U is the number of target tokens for the current timestep Ti.
NOTE: preserve_alignments is an invalid argument for any search_type other than basic beam search.
ngram_lm_model – str The path to the N-gram LM
ngram_lm_alpha – float Alpha weight of N-gram LM
tokens_type – str Tokenization type [‘subword’, ‘char’]
- align_length_sync_decoding(h: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[Hypothesis] = None) List[Hypothesis] [source]#
Alignment-length synchronous beam search implementation. Based on https://ieeexplore.ieee.org/document/9053040
- Parameters
h – Encoded speech features (1, T_max, D_enc)
- Returns
N-best decoding results
- Return type
nbest_hyps
- compute_ngram_score(current_lm_state: kenlm.State, label: int) Tuple[float, kenlm.State] [source]#
Score computation for kenlm ngram language model.
- default_beam_search(h: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[Hypothesis] = None) List[Hypothesis] [source]#
Beam search implementation.
- Parameters
x – Encoded speech features (1, T_max, D_enc)
- Returns
N-best decoding results
- Return type
nbest_hyps
- greedy_search(h: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[Hypothesis] = None) List[Hypothesis] [source]#
Greedy search implementation for transducer. Generic case when beam size = 1. Results might differ slightly due to implementation details as compared to GreedyRNNTInfer and GreedyBatchRNNTInfer.
- Parameters
h – Encoded speech features (1, T_max, D_enc)
- Returns
1-best decoding results
- Return type
hyp
- property input_types#
Returns definitions of module input ports.
- modified_adaptive_expansion_search(h: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[Hypothesis] = None) List[Hypothesis] [source]#
Based on/modified from https://ieeexplore.ieee.org/document/9250505
- Parameters
h – Encoded speech features (1, T_max, D_enc)
- Returns
N-best decoding results
- Return type
nbest_hyps
- property output_types#
Returns definitions of module output ports.
- prefix_search(hypotheses: List[Hypothesis], enc_out: torch.Tensor, prefix_alpha: int) List[Hypothesis] [source]#
Prefix search for NSC and mAES strategies. Based on https://arxiv.org/pdf/1211.3711.pdf
- recombine_hypotheses(hypotheses: List[Hypothesis]) List[Hypothesis] [source]#
Recombine hypotheses with equivalent output sequence.
- Parameters
hypotheses (list) – list of hypotheses
- Returns
list of recombined hypotheses
- Return type
final (list)
- resolve_joint_output(enc_out: torch.Tensor, dec_out: torch.Tensor) Tuple[torch.Tensor, torch.Tensor] [source]#
Resolve output types for RNNT and HAT joint models
- sort_nbest(hyps: List[Hypothesis]) List[Hypothesis] [source]#
Sort hypotheses by score or score given sequence length.
- Parameters
hyps – list of hypotheses
- Returns
sorted list of hypotheses
- Return type
hyps
- time_sync_decoding(h: torch.Tensor, encoded_lengths: torch.Tensor, partial_hypotheses: Optional[Hypothesis] = None) List[Hypothesis] [source]#
Time synchronous beam search implementation. Based on https://ieeexplore.ieee.org/document/9053040
- Parameters
h – Encoded speech features (1, T_max, D_enc)
- Returns
N-best decoding results
- Return type
nbest_hyps
Hypotheses#
- class nemo.collections.asr.parts.utils.rnnt_utils.Hypothesis(score: float, y_sequence: ~typing.Union[~typing.List[int], torch.Tensor], text: ~typing.Optional[str] = None, dec_out: ~typing.Optional[~typing.List[torch.Tensor]] = None, dec_state: ~typing.Optional[~typing.Union[~typing.List[~typing.List[torch.Tensor]], ~typing.List[torch.Tensor]]] = None, timestep: ~typing.Union[~typing.List[int], torch.Tensor] = <factory>, alignments: ~typing.Optional[~typing.Union[~typing.List[int], ~typing.List[~typing.List[int]]]] = None, frame_confidence: ~typing.Optional[~typing.Union[~typing.List[float], ~typing.List[~typing.List[float]]]] = None, token_confidence: ~typing.Optional[~typing.List[float]] = None, word_confidence: ~typing.Optional[~typing.List[float]] = None, length: ~typing.Union[int, torch.Tensor] = 0, y: ~typing.Optional[~typing.List[torch.tensor]] = None, lm_state: ~typing.Optional[~typing.Union[~typing.Dict[str, ~typing.Any], ~typing.List[~typing.Any]]] = None, lm_scores: ~typing.Optional[torch.Tensor] = None, ngram_lm_state: ~typing.Optional[~typing.Union[~typing.Dict[str, ~typing.Any], ~typing.List[~typing.Any]]] = None, tokens: ~typing.Optional[~typing.Union[~typing.List[int], torch.Tensor]] = None, last_token: ~typing.Optional[torch.Tensor] = None)[source]#
Bases:
object
Hypothesis class for beam search algorithms.
score: A float score obtained from an AbstractRNNTDecoder module’s score_hypothesis method.
- y_sequence: Either a sequence of integer ids pointing to some vocabulary, or a packed torch.Tensor
behaving in the same manner. dtype must be torch.Long in the latter case.
dec_state: A list (or list of list) of LSTM-RNN decoder states. Can be None.
- text: (Optional) A decoded string after processing via CTC / RNN-T decoding (removing the CTC/RNNT
blank tokens, and optionally merging word-pieces). Should be used as decoded string for Word Error Rate calculation.
- timestep: (Optional) A list of integer indices representing at which index in the decoding
process did the token appear. Should be of same length as the number of non-blank tokens.
- alignments: (Optional) Represents the CTC / RNNT token alignments as integer tokens along an axis of
time T (for CTC) or Time x Target (TxU). For CTC, represented as a single list of integer indices. For RNNT, represented as a dangling list of list of integer indices. Outer list represents Time dimension (T), inner list represents Target dimension (U). The set of valid indices includes the CTC / RNNT blank token in order to represent alignments.
- frame_confidence: (Optional) Represents the CTC / RNNT per-frame confidence scores as token probabilities
along an axis of time T (for CTC) or Time x Target (TxU). For CTC, represented as a single list of float indices. For RNNT, represented as a dangling list of list of float indices. Outer list represents Time dimension (T), inner list represents Target dimension (U).
- token_confidence: (Optional) Represents the CTC / RNNT per-token confidence scores as token probabilities
along an axis of Target U. Represented as a single list of float indices.
- word_confidence: (Optional) Represents the CTC / RNNT per-word confidence scores as token probabilities
along an axis of Target U. Represented as a single list of float indices.
- length: Represents the length of the sequence (the original length without padding), otherwise
defaults to 0.
y: (Unused) A list of torch.Tensors representing the list of hypotheses.
lm_state: (Unused) A dictionary state cache used by an external Language Model.
lm_scores: (Unused) Score of the external Language Model.
ngram_lm_state: (Optional) State of the external n-gram Language Model.
tokens: (Optional) A list of decoded tokens (can be characters or word-pieces.
last_token (Optional): A token or batch of tokens which was predicted in the last step.
- class nemo.collections.asr.parts.utils.rnnt_utils.NBestHypotheses(n_best_hypotheses: Optional[List[Hypothesis]])[source]#
Bases:
object
List of N best hypotheses
Adapter Networks#
- class nemo.collections.asr.parts.submodules.adapters.multi_head_attention_adapter_module.MultiHeadAttentionAdapter(*args: Any, **kwargs: Any)[source]#
Bases:
MultiHeadAttention
,AdapterModuleUtil
Multi-Head Attention layer of Transformer. :param n_head: number of heads :type n_head: int :param n_feat: size of the features :type n_feat: int :param dropout_rate: dropout rate :type dropout_rate: float :param proj_dim: Optional integer value for projection before computing attention.
If None, then there is no projection (equivalent to proj_dim = n_feat). If > 0, then will project the n_feat to proj_dim before calculating attention. If <0, then will equal n_head, so that each head has a projected dimension of 1.
- forward(query, key, value, mask, pos_emb=None, cache=None)[source]#
Compute ‘Scaled Dot Product Attention’. :param query: (batch, time1, size) :type query: torch.Tensor :param key: (batch, time2, size) :type key: torch.Tensor :param value: (batch, time2, size) :type value: torch.Tensor :param mask: (batch, time1, time2) :type mask: torch.Tensor :param cache: (batch, time_cache, size) :type cache: torch.Tensor
- Returns
transformed value (batch, time1, d_model) weighted by the query dot key attention cache (torch.Tensor) : (batch, time_cache_next, size)
- Return type
output (torch.Tensor)
- class nemo.collections.asr.parts.submodules.adapters.multi_head_attention_adapter_module.RelPositionMultiHeadAttentionAdapter(*args: Any, **kwargs: Any)[source]#
Bases:
RelPositionMultiHeadAttention
,AdapterModuleUtil
Multi-Head Attention layer of Transformer-XL with support of relative positional encoding. Paper: https://arxiv.org/abs/1901.02860 :param n_head: number of heads :type n_head: int :param n_feat: size of the features :type n_feat: int :param dropout_rate: dropout rate :type dropout_rate: float :param proj_dim: Optional integer value for projection before computing attention.
If None, then there is no projection (equivalent to proj_dim = n_feat). If > 0, then will project the n_feat to proj_dim before calculating attention. If <0, then will equal n_head, so that each head has a projected dimension of 1.
- Parameters
adapter_strategy – By default, MHAResidualAddAdapterStrategyConfig. An adapter composition function object.
- forward(query, key, value, mask, pos_emb, cache=None)[source]#
Compute ‘Scaled Dot Product Attention’ with rel. positional encoding. :param query: (batch, time1, size) :type query: torch.Tensor :param key: (batch, time2, size) :type key: torch.Tensor :param value: (batch, time2, size) :type value: torch.Tensor :param mask: (batch, time1, time2) :type mask: torch.Tensor :param pos_emb: (batch, time1, size) :type pos_emb: torch.Tensor :param cache: (batch, time_cache, size) :type cache: torch.Tensor
- Returns
transformed value (batch, time1, d_model) weighted by the query dot key attention cache_next (torch.Tensor) : (batch, time_cache_next, size)
- Return type
output (torch.Tensor)
- class nemo.collections.asr.parts.submodules.adapters.multi_head_attention_adapter_module.PositionalEncodingAdapter(*args: Any, **kwargs: Any)[source]#
Bases:
PositionalEncoding
,AdapterModuleUtil
Absolute positional embedding adapter.
Note
Absolute positional embedding value is added to the input tensor without residual connection ! Therefore, the input is changed, if you only require the positional embedding, drop the returned x !
- Parameters
d_model (int) – The input dimension of x.
max_len (int) – The max sequence length.
xscale (float) – The input scaling factor. Defaults to 1.0.
adapter_strategy (AbstractAdapterStrategy) – By default, ReturnResultAdapterStrategyConfig. An adapter composition function object. NOTE: Since this is a positional encoding, it will not add a residual !
- class nemo.collections.asr.parts.submodules.adapters.multi_head_attention_adapter_module.RelPositionalEncodingAdapter(*args: Any, **kwargs: Any)[source]#
Bases:
RelPositionalEncoding
,AdapterModuleUtil
Relative positional encoding for TransformerXL’s layers See : Appendix B in https://arxiv.org/abs/1901.02860
Note
Relative positional embedding value is not added to the input tensor ! Therefore, the input should be updated changed, if you only require the positional embedding, drop the returned x !
- Parameters
d_model (int) – embedding dim
max_len (int) – maximum input length
xscale (bool) – whether to scale the input by sqrt(d_model)
adapter_strategy – By default, ReturnResultAdapterStrategyConfig. An adapter composition function object.
Adapter Strategies#
- class nemo.collections.asr.parts.submodules.adapters.multi_head_attention_adapter_module.MHAResidualAddAdapterStrategy(stochastic_depth: float = 0.0, l2_lambda: float = 0.0)[source]#
Bases:
ResidualAddAdapterStrategy
An implementation of residual addition of an adapter module with its input for the MHA Adapters.
- forward(input: torch.Tensor, adapter: torch.nn.Module, *, module: AdapterModuleMixin)[source]#
A basic strategy, comprising of a residual connection over the input, after forward pass by the underlying adapter. Additional work is done to pack and unpack the dictionary of inputs and outputs.
Note: The value tensor is added to the output of the attention adapter as the residual connection.
- Parameters
input –
A dictionary of multiple input arguments for the adapter module. query, key, value: Original output tensor of the module, or the output of the
previous adapter (if more than one adapters are enabled). mask: Attention mask. pos_emb: Optional positional embedding for relative encoding.
adapter – The adapter module that is currently required to perform the forward pass.
module – The calling module, in its entirety. It is a module that implements AdapterModuleMixin, therefore the strategy can access all other adapters in this module via module.adapter_layer.
- Returns
The result tensor, after one of the active adapters has finished its forward passes.
- compute_output(input: torch.Tensor, adapter: torch.nn.Module, *, module: AdapterModuleMixin) torch.Tensor [source]#
Compute the output of a single adapter to some input.
- Parameters
input – Original output tensor of the module, or the output of the previous adapter (if more than one adapters are enabled).
adapter – The adapter module that is currently required to perform the forward pass.
module – The calling module, in its entirety. It is a module that implements AdapterModuleMixin, therefore the strategy can access all other adapters in this module via module.adapter_layer.
- Returns
The result tensor, after one of the active adapters has finished its forward passes.