Support dataloader as input to audio for transcription

[titu1994]

What does this PR do ?

Enables the use of a pre-constructed data loader as input to the model.transcribe() function.
This allows for a fastpath to ignore all manifest and tensor handling to the user, only executing the model forward and later steps.

Collection: [ASR]

Changelog

• Allows the user to provide a DataLoader object, which overrides internal computation of manifest processing or dataset construction
• Assumes implicit faith in user provided input - the user is now responsible for formatting and providing all arguments to match up with the ASR model's forward arguments if user chooses to provide a dataloader.

Usage

from nemo.collections.asr.data.audio_to_text import _speech_collate_fn

model = ASRModel.from_pretrained("stt_en_conformer_ctc_small")

# Load audio file
import soundfile as SF

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_file2 = os.path.join(test_data_dir, "asr", "train", "an4", "wav", "an152-mwhw-b.wav")
audio2, sr = sf.read(audio_file2, dtype='float32')

# Create a dummy dataset to hold the tensor values
class DummyDataset(Dataset):
    def __init__(self, audio_tensors: List[str], config: Dict = None):
        self.audio_tensors = audio_tensors
        self.config = config

    def __getitem__(self, index):
        data = self.audio_tensors[index]
        samples = torch.tensor(data)
        # Calculate seq length
        seq_len = torch.tensor(samples.shape[0], dtype=torch.long)

        # Dummy text tokens
        text_tokens = torch.tensor([0], dtype=torch.long)
        text_tokens_len = torch.tensor(1, dtype=torch.long)

        # Ensure to provide output tokens that can be consumed by an ASR's forward function
        return (samples, seq_len, text_tokens, text_tokens_len)

    def __len__(self):
        return len(self.audio_tensors)

# Wrap the dataset into a data loader with proper collate function
dataset = DummyDataset([audio, audio2])
collate_fn = lambda x: _speech_collate_fn(x, pad_id=0)
dataloader = DataLoader(dataset, batch_size=2, shuffle=False, num_workers=0, collate_fn=collate_fn)

# DataLoader as input to audio
outputs = model.transcribe(dataloader, batch_size=1)

assert len(outputs) == 2
assert isinstance(outputs[0], str)
assert isinstance(outputs[1], str)

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

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

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[titu1994]

Fyi @nithinraok - this should enable fastpath execution for HF too

[titu1994]

@galv @pzelasko actual change is just these 5 lines, everything else is just black formatter and a unittest at the bottom - https://github.com/NVIDIA/NeMo/pull/9201/files#diff-04e10f8fb8f7afddb360c7ea0ff9c831613c754085d2fd06961bb80a9f932a25R372-R376

[pzelasko]

LGTM.

@titu1994 @nithinraok @galv I'd like to point out that if you're going to use a TensorDataset, make it an IterableDataset that yields collated mini-batches, rather than relying on DataLoader's batch_size and collate_fn to do it for you. This way you'll amortize the overhead of collation as it will happen in the background process. This change will matter for super-high RTFx models.

[nithinraok]

minor comments. For huggingface, we need to depend on Iterable dataset I believe.

[nithinraok]

Add torch.dataloader to audio types in the func signature.

[titu1994]

Done

[nithinraok]

Add torch.dataloader to audio types in the func signature.

[titu1994]

Done

[galv]

@nithinraok can you investigate to see whether this change will fix your performance degredation issues when using transcribe() within the huggingface open asr leaderboard RTFx measurements?

[nithinraok]

@galv yeah, I am thinking to use iterable dataset (looks like streaming is the option for huggingface datasets to get iterable dataset) and run the HF evals.