I'm hitting this error message now and then. It does not seem to be affecting training, but I only see it when training on TPU. The same dataset was used in GPU with no errors. Just posting here in case there is something else going on that I am missing.

Step... (75000/759160 | Eval Loss: 0.11205478757619858 | Eval wer: 0.09877239458498131 | Eval cer: 0.02933955305671511 |):  10%|████▏                                     | 4/40 [23:39:27<205:34:14, 20557.06s/it/
data/flax/lib/python3.8/site-packages/transformers/tokenization_utils_base.py:719: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndar
rays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.                                                                          
  tensor = as_tensor(value)                                                                                                                                                                                        
--- Logging error ---                                                                                                                                                                                              
Traceback (most recent call last):                                                                                                                                                                                 
  File "run_flax_speech_recognition_ctc.py", line 1631, in <module>                                                                                                                                                
    main()                                                                                                                                                                                                         
  File "run_flax_speech_recognition_ctc.py", line 1544, in main                                                                                                                                                    
    state, train_metric = p_train_step(state, batch)                                                                                                                                                               
  File "/data/flax/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback                                                                                           
    return fun(*args, **kwargs)                                                                                                                                                                                    
  File "/data/flax/lib/python3.8/site-packages/jax/_src/api.py", line 2158, in cache_miss                                                                                                                          
    out_tree, out_flat = f_pmapped_(*args, **kwargs)                                                                                                                                                               
  File "/data/flax/lib/python3.8/site-packages/jax/_src/api.py", line 2031, in pmap_f                                                                                                                              
    p = _prepare_pmap(                                                                                                                                                                                             
  File "/data/flax/lib/python3.8/site-packages/jax/_src/api.py", line 1969, in _prepare_pmap                                                                                                                       
    _check_arg(arg)                                                                                                                                                                                                
  File "/data/flax/lib/python3.8/site-packages/jax/_src/api.py", line 2994, in _check_arg                                                                                                                          
    raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")                                                                                                                              
jax._src.traceback_util.UnfilteredStackTrace: TypeError: Argument '[[-100 -100]                                                                                                                                    
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]                                                                                                                                                                                                       
 [-100 -100]]' of type <class 'numpy.ndarray'> is not a valid JAX type.                                                                                                                                            
                                                                                                                                                                                                                   
The stack trace below excludes JAX-internal frames.                                                                                                                                                                
The preceding is the original exception that occurred, unmodified.                                                                                                                                                 
                                                                                                                                                                                                                   
--------------------                                                                                                                                                                                
                                                                                                                                                                                                                   
The above exception was the direct cause of the following exception:                                                                                                                                               
                                                                                                                                                                                                                   
Traceback (most recent call last):                                                                                                                                                                                 
  File "run_flax_speech_recognition_ctc.py", line 1544, in main                                                                                                                                                    
    state, train_metric = p_train_step(state, batch) 
TypeError: Argument '[[-100 -100]
 [-100 -100]
 [-100 -100]
 [-100 -100]
 [-100 -100]
 [-100 -100]
 [-100 -100]
 [-100 -100]]' of type <class 'numpy.ndarray'> is not a valid JAX type.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.8/logging/__init__.py", line 1085, in emit
    msg = self.format(record)
  File "/usr/lib/python3.8/logging/__init__.py", line 929, in format
    return fmt.format(record)
  File "/usr/lib/python3.8/logging/__init__.py", line 668, in format
    record.message = record.getMessage()
  File "/usr/lib/python3.8/logging/__init__.py", line 373, in getMessage
    msg = msg % self.args
TypeError: not all arguments converted during string formatting
Call stack:
  File "run_flax_speech_recognition_ctc.py", line 1631, in <module>
    main()
  File "run_flax_speech_recognition_ctc.py", line 1546, in main
    logger.warning("Encountered following error: \n", e)
Message: 'Encountered following error: \n'
Arguments: (TypeError("Argument '[[-100 -100]\n [-100 -100]\n [-100 -100]\n [-100 -100]\n [-100 -100]\n [-100 -100]\n [-100 -100]\n [-100 -100]]' of type <class 'numpy.ndarray'> is not a valid JAX type."),)                                                  

Training the baseline CTC model on the Common Voice 9 (CV9) dataset, we observe that the training loss drops below zero after ~1.5k train steps: https://wandb.ai/sanchit-gandhi/commonvoice_9_0/runs/y593pwm4?workspace=user-sanchit-gandhi. The CTC loss should be strictly nonnegative.

• CV9 tokenizer: working as expected. Tested within the training script (both tokenising, and decoding), and checked that all attributes all set correctly. Furthermore, the target string in the wandb predictions logs are identical to the transcribed text in the training data -> the tokenizer is correctly tokenising and decoding.
• Logits test: for the randomly initialised (unscanned) model, the PT-Flax equivalence test passes on CV9 for both the logits and losses. Loss is nonnegative. https://github.com/sanchit-gandhi/seq2seq-speech/blob/main/tests/check_flax_ctc_cv9.py
• Trained Flax model: using the 50k train steps checkpoint, the loss is negative.
• We expect to see all log probabilities in the CTC loss function be strictly negative: the probabilities should lie in the range 0 to 1, and so should have a max value of log(1) = 0.
• The CTC loss function can be divided into three stages:

1. Initialisation of log prob arrays
2. Looping over the CTC Markov chain process
3. Extraction of per-sequence loss and CTC reduction ("mean" reduction)

• Our attention is focused on the CTC Markov chain process: https://github.com/sanchit-gandhi/seq2seq-speech/blob/74c4e3650ce1865a61dbf5e99c199539d6815bfc/run_flax_speech_recognition_ctc.py#L638-L659
• When we remove the lax backend in the CTC loss function and use some print statements, we see that positive values creep into the log probabilities by the 4th loop of the CTC algorithm. They first occur in the 'phi-to-emit transition' with the 'next_emit' log probabilities, and then cascade to all further log probs. https://github.com/sanchit-gandhi/seq2seq-speech/blob/main/tests/check_negative_loss_ctc.ipynb
• Changing the value of the log_epsilon hyperparameter to a more negative value does not alter this behaviour.

It seems that all the data are filtered by the input and label lengths in the current implementation, including eval samples, as in the code.

Half of the eval samples for TEDLIUMv3 are not included in the evaluation as a result.

Adds the RNN-T BPE model from NVIDIA NeMo and a training script to train the model with HF Dataset & Trainer. Template training scripts detail the configuration to train a ContextNet ASR model, large size (~144M) with Transducer loss and sub-word encoding.

This issue explores the behaviour of the example CV9 tokenizer created in the get_ctc_tokenizer.py file saved at https://huggingface.co/patrickvonplaten/wav2vec2_ctc_cv9_tokenizer.

With the do_lower_case argument set to true, we observe that the tokenizer returns the <unk> token for all input chars:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/wav2vec2_ctc_cv9_tokenizer", do_lower_case=True)

input_str = "The cat sat on the mat."
input_ids = tokenizer(input_str).input_ids
decoded_str = tokenizer.decode(input_ids)

print("Input str: ", input_str)
print("Input ids: ", input_ids)
print("Decoded str: ", decoded_str)
print("<unk> token id: ", tokenizer.unk_token_id)
print("Type: ", type(tokenizer))

Output:

Input str:  The cat sat on the mat.
Input ids:  [3, 3, 3, 25, 3, 3, 3, 25, 3, 3, 3, 25, 3, 3, 25, 3, 3, 3, 25, 3, 3, 3, 5]
Decoded str:  <unk> <unk> <unk> <unk> <unk> <unk>.
<unk> token id:  3
Type:  <class 'transformers.models.wav2vec2.tokenization_wav2vec2.Wav2Vec2CTCTokenizer'>

Every input character in the alphabet is incorrectly mapped to the unknown token. However, what's interesting here is that we see the punctuation (word-spaces and full-stop) correctly tokenized and decoded. These are the two inputs that are unaffected by the do_lower_case operation.

If we now re-run this operation with the do_lower_case argument set to False, we observe that the decoded word string is correct (barring the first capital T):

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/wav2vec2_ctc_cv9_tokenizer", do_lower_case=False)

input_str = "The cat sat on the mat."
input_ids = tokenizer(input_str).input_ids
decoded_str = tokenizer.decode(input_ids)

print("Input str: ", input_str)
print("Input ids: ", input_ids)
print("Decoded str: ", decoded_str)

Output:

Input str:  The cat sat on the mat.
Input ids:  [3, 9, 7, 25, 17, 35, 18, 25, 8, 35, 18, 25, 16, 6, 25, 18, 9, 7, 25, 37, 35, 18, 5]
Decoded str:  <unk>he cat sat on the mat.

Thus, the issue must lie within the do_lower_case attribute of the Wav2Vec2CTCTokenizer. Indeed, if we inspect the modelling code for this tokenizer, we observe that setting do_lower_case to True first runs the upper method on the input string, converting the string to upper-case: https://github.com/huggingface/transformers/blob/6d80c92c77593dc674052b5a46431902e6adfe88/src/transformers/models/wav2vec2/tokenization_wav2vec2.py#L238 Since the tokenizer was built on the lower-case vocabulary, each of these upper-case characters are OOV, and so assigned to the <unk> token.

What we should do is set the do_lower_case argument to False, meaning the upper-case operation is not performed on the input string, and simply lower-case the input string before tokenizing:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/wav2vec2_ctc_cv9_tokenizer", do_lower_case=False)

input_str = "The cat sat on the mat.".lower()
input_ids = tokenizer(input_str).input_ids
decoded_str = tokenizer.decode(input_ids)

print("Input str: ", input_str)
print("Input ids: ", input_ids)
print("Decoded str: ", decoded_str)

Input str:  the cat sat on the mat.
Input ids:  [18, 9, 7, 25, 17, 35, 18, 25, 8, 35, 18, 25, 16, 6, 25, 18, 9, 7, 25, 37, 35, 18, 5]
Decoded str:  the cat sat on the mat.

This simply required modification of the lines which instantiate the tokenizer from pretrained in the CTC training script: https://github.com/sanchit-gandhi/seq2seq-speech/blob/de22a44b4e902a7ccb185f7f1f44053bc3ad297f/run_flax_speech_recognition_ctc.py#L809-L816

We should only filter training samples by our audio and transcription length criterions when fine-tuning ASR systems. The eval and test sets should not be filtered. Doing so leads to partial datasets, and thus in-valid results.

Perform error correction of datasets on a dataset-by-dataset basis:

LibriSpeech ASR:

• No error correction necessary

VoxPopuli:

• No error correction necessary

Common Voice 9:

• Remove trailing quotation marks as they do not affect the transcription
• Replace double quotation marks with single
• Normalise quotation marks, apostrophes, hyphens
• Normalise the text to always finish with punctuation

TED-LIUM:

• Remove tokens
• Replace spaced apostrophes by un-spaced for most frequent instances (it 's -> it's)*

GigaSpeech:

• Remove junk tokens (disfluencies)
• Convert spelled out punctuation to symbolic ( -> ,)
• Normalise the text to always finish with punctuation

SwitchBoard:

• Remove junk tokens (disfluencies)
• Remove parenthesised text detailing annotations (the cat sat on the (2) mat -> the cat sat on the mat)
• Replace anomalous words with their correct transcriptions: the cat/ca sat on the/th mat -> the cat sat on the mat

Earnings22:

• Remove junk tokens (disfluencies)
• Replace mal-formatted ellipsis with spaced (… -> . . . )

JIWER compliance (for WER/CER metrics):

• Remove multiple spaces
• Strip trailing spaces

Contraction Occurrences
Total contractions in dataset: 120399
-----------------------------------------------------
     Cont. |       Total occ |       % of total occ |
-----------------------------------------------------
        's |           56216 |               46.691 |
        't |           22628 |               18.794 |
       're |           16658 |               13.836 |
       've |            9464 |                7.861 |
        'm |            8360 |                6.944 |
       'll |            3636 |                 3.02 |
        'd |            3259 |                2.707 |
    'clock |              75 |                0.062 |
      'all |              15 |                0.012 |
       'am |               7 |                0.006 |
   'ivoire |               4 |                0.003 |
    'grady |               4 |                0.003 |
    'neill |               4 |                0.003 |
    'brien |               3 |                0.002 |
     'oeil |               3 |                0.002 |
        'n |               3 |                0.002 |
    'alene |               3 |                0.002 |
    'arche |               2 |                0.002 |
        'a |               2 |                0.002 |
      'ivo |               2 |                0.002 |
      'ren |               2 |                0.002 |
     'ites |               2 |                0.002 |
   'connor |               2 |                0.002 |
       'an |               2 |                0.002 |
       'or |               2 |                0.002 |
       'ts |               2 |                0.002 |
     'dour |               2 |                0.002 |
      'tat |               2 |                0.002 |
     'hara |               2 |                0.002 |
   'aquila |               2 |                0.002 |
    'oreal |               1 |                0.001 |
     'oral |               1 |                0.001 |
   'agnese |               1 |                0.001 |
      'tre |               1 |                0.001 |
      'arc |               1 |                0.001 |
     'neal |               1 |                0.001 |
     'shea |               1 |                0.001 |
       'in |               1 |                0.001 |
      'lin |               1 |                0.001 |
  'gieblyn |               1 |                0.001 |
       'ed |               1 |                0.001 |
       'on |               1 |                0.001 |
        'i |               1 |                0.001 |
 'ospedale |               1 |                0.001 |
  'addaura |               1 |                0.001 |
   'reilly |               1 |                0.001 |
  'donohue |               1 |                0.001 |
 'donoghue |               1 |                0.001 |
      'ile |               1 |                0.001 |
  'foscari |               1 |                0.001 |
'herpiniere |               1 |                0.001 |
      'est |               1 |                0.001 |
   'conner |               1 |                0.001 |
  'oeuvres |               1 |                0.001 |
 'academie |               1 |                0.001 |
    'brian |               1 |                0.001 |
  'connell |               1 |                0.001 |
    'byrne |               1 |                0.001 |
   'angelo |               1 |                0.001 |
      'ant |               1 |                0.001 |
      'mun |               1 |                0.001 |

Hi, thanks for these collection of scripts!

I've been trying to run your run_flax_speech_recognition_ctc.py on a single TPUv3-8 but after a few epochs I tend to always run out of memory (not sure if caused by memory leak or something). I also tried to recover from the last checkpoint by skipping the number of steps the model was last saved at, and setting the learning rate appropriately. I also tried modifying MixedPrecisionTrainState.create() to it starts at the last saved checkpoint step too. Nothing worked. As soon as it starts training, it runs out of memory again. Any idea of what could be happening?

Thanks!