Online Code Switching Dataset for ASR

[trias702]

What does this PR do ?

Adds an IterableDataset class which generates on-the-fly code switched utterances for ASR

Collection: ASR

Changelog

• nemo/collections/asr/data/audio_to_text.py: added language support to the BPE classes to support code-switching

• nemo/collections/asr/data/audio_to_text_dataset.py: added the core create_cs_dataset function and various instantiation logic

• nemo/collections/common/data/__init__.py: exposed new CodeSwitchedDataset class

• nemo/collections/common/data/dataset.py: contains new CodeSwitchedDataset class

• nemo/collections/common/parts/preprocessing/collections.py: added discrete language support of the TokeniserWrapper via a hasattr check

• nemo/collections/asr/models/*: changed all of the models to check for IterableDataset instead of reading the is_tarred config item, as I think this makes it more robust

Usage

There are many different permutations to using this new dataset, but a basic example shows how to add it to an existing training config yaml. To understand what each parameter does, please look at the class doc for the CodeSwitchedDataset class in nemo.collections.common.data.dataset.

++model.train_ds.manifest_filepath=[/path/to/en/tarred_manifest.json,/path/to/de/tarred_manifest.json,/path/to/ja/tarred_manifest.json]
++model.train_ds.tarred_audio_filepaths=['/path/to/en/tars/audio__OP_0..511_CL_.tar','/path/to/de/tars/audio__OP_0..1023_CL_.tar', '/path/to/ja/tars/audio__OP_0..2047_CL_.tar']
++model.train_ds.is_code_switched=true
++model.train_ds.is_tarred=true
++model.train_ds.shuffle=true
++model.train_ds.code_switched.min_duration=12
++model.train_ds.code_switched.max_duration=20
++model.train_ds.code_switched.min_monolingual=0.3
++model.train_ds.code_switched.probs=[0.25, 0.5, 0.25]
++model.train_ds.code_switched.force_monochannel=true
++model.train_ds.code_switched.sampling_scales=[0.75, 1, 0.25]
++model.validation_ds.manifest_filepath=[[/path/to/en/non-tarred/manifest.json, /path/to/de/non-tarred/manifest.json, /path/to/ja/non-tarred/manifest.json]]
++model.validation_ds.is_code_switched=true
++model.validation_ds.is_tarred=false \
++model.validation_ds.code_switched.min_duration=6
++model.validation_ds.code_switched.max_duration=20
++model.validation_ds.code_switched.min_monolingual=0.3
++model.validation_ds.code_switched.probs=[0.33, 0.34, 0.33]
++model.validation_ds.code_switched.force_monochannel=true
++model.validation_ds.code_switched.sampling_scales=1
~model.tokenizer
++model.tokenizer.type=agg
++model.tokenizer.langs.en.dir=/path/to/en/tokeniser
++model.tokenizer.langs.en.type=bpe
++model.tokenizer.langs.de.dir=/path/to/de/tokeniser
++model.tokenizer.langs.de.type=bpe
++model.tokenizer.langs.ja.dir=/path/to/ja/tokeniser
++model.tokenizer.langs.ja.type=bpe

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)

[github-actions]

This PR is stale because it has been open for 14 days with no activity. Remove stale label or comment or update or this will be closed in 7 days.

[trias702]

Can someone please review this PR for merging? We really need this functionality for Riva work.

[KunalDhawan]

Great work @trias702! Added a few minor comments but overall the implementation looks good. @bmwshop for final review.

[titu1994]

Overall, the pr needs some documentation (inline as well as in NeMo Docs) and checks.
Core logic of dataset, ill leave to others to review

[titu1994]

This will be super slow right? Better to just not allow augmentation here then?

[trias702]

So this is the part where we apply the augmentor to the final, synthetic utterance only, and why we always pass Augmentor=None to the individual language sub-datasets. We want all the individual samples from the individual mono-languages to be clean, and then we build the clean synthetic, and then we hit that final synthetic with the augmentation.

As for the logic, it's all in-memory operations with io.BytesIO, nothing is actually being written to disk. I have done several training runs with online CS and augmentation, and it doesn't slow it down by much at all, still very much usable.

There's sadly no way around it, we need to do augmentation on the final, full synthetic sample because many augmentations depend on features of the full waveform to work correctly (for example the new NormPerturbation).

[bmwshop]

What happens if the a manifest contains LID information and we pass LIDs here as well?

[trias702]

It's not a problem, because LIDs in the manifest always take precedence, you can see that here:

NeMo/nemo/collections/common/parts/preprocessing/collections.py

Line 160 in 40dd53d

if lang is not None:

The new language parameter is only used as a fallback IFF using AggTokeniser and the manifest has no LID in the manifest.

[KunalDhawan]

Thanks @trias702, LGTM!