Improve mAES algorithm with patches

[titu1994]

What does this PR do ?

Adds the fixes from espnet/espnet#4312 for mAES RNNT beam search.

Collection: [ASR]

Changelog

• Add deduplication checks
• Improve speed of mAES beam search

Usage

# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytorch_lightning as pl
from omegaconf import OmegaConf, open_dict

from nemo.collections.asr.models import EncDecRNNTBPEModel

def main():

    trainer = pl.Trainer(gpus=[1], precision=16, logger=False)

    # Model
    # path = 'some model path'

    asr_model = EncDecRNNTBPEModel.restore_from(path, trainer=trainer, map_location='cpu')

    asr_model.eval()
    asr_model.freeze()

    asr_model.preprocessor.featurizer.dither = 0.0
    asr_model.preprocessor.featurizer.pad_to = 0

    cfg = asr_model.cfg

    test_cfg = cfg.test_ds
    test_cfg.manifest_filepath = "manifests/test_other.json"
    test_cfg.batch_size = 256

    asr_model.setup_multiple_test_data(test_cfg)

    decoding_cfg = cfg.decoding

    with open_dict(decoding_cfg):
        decoding_cfg.strategy = "maes"
        decoding_cfg.fused_batch_size = -1
        decoding_cfg.beam.beam_size = 4
        decoding_cfg.beam.tsd_max_sym_exp = 10
        decoding_cfg.beam.alsd_max_target_len = 2.5
        decoding_cfg.beam.maes_num_steps = 2
        decoding_cfg.beam.maes_prefix_alpha = 1
        decoding_cfg.beam.maes_expansion_beta = 2
        decoding_cfg.beam.maes_expansion_gamma = 2.3

    asr_model.change_decoding_strategy(decoding_cfg)

    if asr_model.prepare_test(trainer):
        trainer.test(asr_model)


if __name__ == '__main__':
    main()

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

[VahidooX]

LGTM, just a minor comment.

[VahidooX]

Should we drop this line?

[titu1994]

Ah i kinda kept it cause i might need it for debugging during rnnt beam setup