ASR Transcription Refactor
State: closed | Pull request: yes
Labels: name: ASR

# What does this PR do ?

Refactors the Transcribe API to be more configurable and reduce the duplication throughout multiple models

Note: The spec is optional to some classes since they are higher level models that operate on entire NeMo Models and as such would require a complete override of all aspects of this API.

**Collection**: [ASR]

# Changelog 
- New `TranscriptionMixin` mixin class that holds the responsibility for providing unified transcription semantics to all classes that implement its abstract methods.
- Removal of `transcribe()` as a function of `ASRModel`. 
  - Reason: ASRModel has become a defacto base class for all models under the speech to text paradigm, which includes classification models, diarizers, SLU, and AST. However the term "transcribe" usually is reserved for basic ASR
  - Subclasses must now chose to opt into transcription via the `TranscriptionMixin`
- Refactor all major classes that utilize the transcription API
- Add support for tensor (numpy ndarray, torch tensor) as input to transcribe() function. Can also accept list of numpy arrays to perform batch inference.

# Breaking Changes
The API mostly tries to be backward compatible, with the following exceptions:
- The first argument is no longer called `path2audio_files` and is now instead just `audio`. This may break cases where keyword argument is used to call `model.transcribe()`. This is because this API aims to unify audio filepath via list and pytorch / numpy tensor as input for single file transcription. 
- Deprecation and removal of `logprobs` argument. You must use `return_hypothesis` argument now and access the logprobs via the `alignments` variable of the Hypothesis dataclass.


# Usage

## General Transcription
`model.transcribe()` remains the general API after refactor.

## Generator-based transcription
Additionall you can now treat the transcription engine as a generator (though you must explicitly pass a dataclass for the config of the transition engine in this case).

This might be useful in cases where lists of hundreds of thousands of samples are provided to the `audio` argument. Previously, all results would be hosted on CPU RAM (including full logprobs or hypothesis objects) and if the ram is insufficient it would OOM and the process would be killed. 

Generator-based inference will make this a constant memory cost, both for the GPU and CPU.

```python
import nemo.collections.asr as nemo_asr
model = nemo_asr.models.ASRModel.from_pretrained("nvidia/stt_en_conformer_transducer_large")

config = nemo_asr.models.TranscribeConfig(batch_size=32)
generator = model.transcribe_generator(audio, config)

for processed_outputs in generator:
    # process a batch of 32 results (or less if last batch does not contain 32 elements)
    ....
```

## Audio tensor inference

NeMo ASR now supports numpy or torch tensors as input to the transcribe function (also includes the ability to use list of tensors for batched inference). 

Note: All numpy tensors are cast to torch.Tensor() (CPU) internally, then moved automatically to device of the model that called transcribe(). Movement to correct device is automatic for both numpy and pytorch tensors.

The API is uniform with filepath based inference:

```python
import torch
import soundfile as sf

from nemo.collections.asr.models import ASRModel
model = ASRModel.from_pretrained("stt_en_conformer_ctc_small")
model.eval()

# Load audio file
audio_file = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an46-mmap-b.wav")
audio, sr = sf.read(audio_file, dtype='float32')

audio_file_2 = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an104-mrcb-b.wav")
audio_2, sr = sf.read(audio_file_2, dtype='float32')
 
# Mix second audio to torch.tensor()
audio_2 = torch.from_numpy(audio_2)

# Numpy array + torch tensor mixed tensor input (for batched inference)
outputs = model.transcribe([audio, audio_2], batch_size=2)
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [x] 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**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation


jenkins

jenkins

jenkins

jenkins

jenkins