TarredAudioToTextDataset assumes audio is always wav
State: closed | Pull request: no

**Describe the bug**

`TarredAudioToTextDataset` assumes audio is `wav`. However I was able to train a model with `mp3` files and would like `TarredAudioToTextDataset` to support other file types.

**Steps/Code to reproduce bug**

Use any dataset with non-wav files.

**Expected behavior**

Model should be trainable with non-wav tarred datasets

**Environment overview (please complete the following information)**

 - Environment location: Bare-metal
 - Method of NeMo install: pip install nemo['all']==1.0.0rc1

**Environment details**

Otherwise, please provide:
- OS version: Ubuntu 20.04
- PyTorch version: 1.7.1
- Python version: 3.8

**Additional context**

My suggestion is to simply add another field to the `*_ds` config called `file_format` or `file_extension` which can be one of a list of supported files by `soundfile`.
Can submit a PR if this is fine, but some guidance/feedback on the approach first would be good.


We generally don't prefer to have other file formats for efficiency reasons (wav takes storage but at least it prevents on the fly resampling to 16khz monochannel audio). 

Tarred datasets are meant for distributed training on clusters, where network I/O can have a significant impact on training speed. Likewise, it would be unwise to spend compute for such expensive resampling on the CPU repeatedly per sample.

In any case, I'll look into what's currently preventing Tarred datasets of ASR from working with other file formats. 

I should've left a more detailed trace, sorry.
The issue happens [here](https://github.com/NVIDIA/NeMo/blob/489bb241a6ed727b4d7ad3213242772a7628ff55/nemo/collections/asr/data/audio_to_text.py#L920) where the `webdataset` try's to rename the audio key from `wav`.

Instead, this should be an arg passed to `_TarredAudioToTextDataset`.

Otherwise, you get the following trace 

```
Original Traceback (most recent call last):                                                                            [1447/1909]
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/utils/data/_utils/worker.py", line 198, in _worker_loop             
    data = fetcher.fetch(index)                                                                                                   
  File "/home/ubuntu/.local/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 28, in fetch                      
    data.append(next(self.dataset_iter))                                                                                          
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 226, in map                                
    for sample in data:                                                                                                           
  File "/home/ubuntu/.local/lib/python3.8/site-packages/nemo/collections/asr/data/audio_to_text.py", line 754, in __next__        
    audio_bytes, audio_filename = next(self.iterator)                                                                             
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 282, in to_tuple                           
    for sample in data:                                                                                                           
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 244, in rename                             
    if handler(exn):                                                                                                              
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/utils.py", line 12, in reraise_exception                       
    raise exn                                                                                                                     
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 242, in rename                             
    yield {k: getfirst(sample, v, missing_is_error=True) for k, v in kw.items()}                                                  
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 242, in <dictcomp>                         
    yield {k: getfirst(sample, v, missing_is_error=True) for k, v in kw.items()}                                                  
  File "/home/ubuntu/.local/lib/python3.8/site-packages/webdataset/iterators.py", line 77, in getfirst                            
    raise ValueError(f"didn't find {keys} in {list(a.keys())}")                                                                   
ValueError: didn't find ['wav'] in ['__key__', 'mp3']                                                                             
```      

That makes sense, thanks for the hint!