Extended input configuration + Lhotse multimodal (mixed audio and text-only) dataloading
State: closed | Pull request: yes
Labels: name: ASR

name: common

# What does this PR do ?

We add a new (optional) way to specify training data for NeMo speech collections. This is intended for experiments which combine potentially many datasets. The configs below are intended as a key under `model.train_ds` in training script configs. We also support providing a path to YAML file with the same contents. 

In addition, we extend Lhotse dataloading support to text-only data (e.g. language modeling or machine translation). For now we assume simple text files, and additional parsers can be easily added in the future. This uses Lhotse version 1.22.0 which introduced support for customizing the samplers to other data than audio. In these cases, we measure the "example length" using the number of tokens for text and define a "token_equivalent_duration" for converting audio duration to a number of "tokens". This enables sampling both modalities together, with dynamic mixing, weighting, dynamic bucketing support, tarred audio, and everything else lhotse provides.

These changes are fully backward compatible (i.e., the new format is optional).

A few examples below:

Example 1. Combine two datasets with equal weights and attach custom metadata in ``tags`` to each cut:

```yaml
input_cfg:
  - type: nemo_tarred
    manifest_filepath: /path/to/manifest__OP_0..512_CL_.json
    tarred_audio_filepath: /path/to/tarred_audio/audio__OP_0..512_CL_.tar
    weight: 0.5
    tags:
      lang: en
      some_metadata: some_value
  - type: nemo_tarred
    manifest_filepath: /path/to/manifest__OP_0..512_CL_.json
    tarred_audio_filepath: /path/to/tarred_audio/audio__OP_0..512_CL_.tar
    weight: 0.5
    tags:
      lang: pl
      some_metadata: some_value
```

(The rest of the config goes in as usual as keys next to `input_cfg`, e.g. `batch_size`, `shuffle`, etc.)

Example 2. Combine multiple (4) datasets, with 2 corresponding to different tasks (ASR, AST). There are two levels of weights: per task (outer) and per dataset (inner). The final weight is the product of outer and inner weight. The groups could be nested further as needed.

```yaml
input_cfg:
  - type: group
    weight: 0.7
    tags:
      task: asr
    input_cfg:
      - type: nemo_tarred
        manifest_filepath: /path/to/asr1/manifest__OP_0..512_CL_.json
        tarred_audio_filepath: /path/to/tarred_audio/asr1/audio__OP_0..512_CL_.tar
        weight: 0.6
        tags:
          lang: en
          some_metadata: some_value
      - type: nemo_tarred
        manifest_filepath: /path/to/asr2/manifest__OP_0..512_CL_.json
        tarred_audio_filepath: /path/to/asr2/tarred_audio/audio__OP_0..512_CL_.tar
        weight: 0.4
        tags:
          lang: pl
          some_metadata: some_value
  - type: group
    weight: 0.3
    tags:
      task: ast
    input_cfg:
      - type: nemo_tarred
        manifest_filepath: /path/to/ast1/manifest__OP_0..512_CL_.json
        tarred_audio_filepath: /path/to/ast1/tarred_audio/audio__OP_0..512_CL_.tar
        weight: 0.2
        tags:
          src_lang: en
          tgt_lang: pl
      - type: nemo_tarred
        manifest_filepath: /path/to/ast2/manifest__OP_0..512_CL_.json
        tarred_audio_filepath: /path/to/ast2/tarred_audio/audio__OP_0..512_CL_.tar
        weight: 0.8
        tags:
          src_lang: pl
          tgt_lang: en
```

Example 3. Combine an ASR (audio-text) dataset with an MT (text-only) dataset so that mini-batches have some examples from both datasets. Provide a custom prompt field for both datasets (to be leveraged by a relevant dataset class):

```yaml
use_multimodal_sampling: true
batch_tokens: 1024
token_equivalent_duration: 0.08  # 0.01 frame shift * 8 subsampling factor
quadratic_factor: 50
num_buckets: 30
use_bucketing: true
input_cfg:
  - type: nemo_tarred
    manifest_filepath: /path/to/manifest__OP_0..512_CL_.json
    tarred_audio_filepath: /path/to/tarred_audio/audio__OP_0..512_CL_.tar
    weight: 0.5
    tags:
      lang: en
      prompt: "Given the following recording, transcribe what the person is saying:"
  - type: txt_pair
    source_path: /path/to/en.txt
    target_path: /path/to/pl.txt
    source_language: en
    target_language: pl
    weight: 0.5
    tags:
      prompt: "Translate the following text to Polish:"
      some_metadata: some_value
```

**Collection**: All speech collections

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# Jenkins CI
To run Jenkins, a NeMo User with write access must comment `jenkins` on the PR.

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] 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

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](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)


jenkins

I'm not convinced that users will be able to jump to this config as easily as Piotr believes, but I won't block it since it's optional and doesn't hurt to have as an option 

> I'm not convinced that users will be able to jump to this config as easily as Piotr believes, but I won't block it since it's optional and doesn't hurt to have as an option

It's definitely an advanced-level user feature for those who need a fine-grained control in complex multidataset/multimodal setups. Hopefully with time we'll be able to release more models and materials demonstrating the usage.

jenkins

jenkins

jenkins