[ASR] Beam Search Decoding How to Get Beam of Tokens as Output
State: closed | Pull request: no

I am using a Conformer-Transducer model and using the `transcribe_speech.py` script for transcript generation.
I am looking to get a time-aligned word-level transcript output, but instead of just the best hypothesis, I want to get N best tokens for each token location, i.e. for each word location, I want the N best predicted words and the time for that location (beam size = N).

For example, with a beam size of 2:

Spoken text: "I can see clearly"

What I have now:
"timestamps_word": [{"word": "I", "start_offset": 16, "end_offset": 20}, {"word": "can", "start_offset": 20, "end_offset": 159}, {"word": "see", "start_offset": 159, "end_offset": 164}, {"word": "clearly", "start_offset": 164, "end_offset": 326}]

Theoretically what I am looking for: "timestamps_word":
[{"words": ["I", "eye"], "start_offset": 16, "end_offset": 20}, {"words": ["can", "cam"], "start_offset": 20, "end_offset": 159}, {"words": ["see", "sea"], "start_offset": 159, "end_offset": 164}, {"words": ["clearly", "freely"], "start_offset": 164, "end_offset": 326}]

My current decoder-specific arguments to the script are as follows:
```decoder_type='rnnt' compute_timestamps=True rnnt_decoding.strategy='beam' rnnt_decoding.beam.beam_size=3 rnnt_decoding.preserve_alignments=True rnnt_decoding.beam.return_best_hypothesis=False```

Is what I am looking for possible?

I don't believe we've implemented beam search for transducer plus greedy word timestamp extraction.

Would this be possible using CTC decoding, if I used a Conformer-CTC model instead?

@titu1994 does this PR not give the results that I am looking for? Or does the 'preserve_alignments' part refer to something else? https://github.com/NVIDIA/NeMo/pull/5925 

Alignment extraction is one required step of greedy word timestamp collection, since that alignment information is used to build the consecutive chunks of sequences that correspond to tokens. 

However while alignment extraction may be implemented for all models, word timestamp calculation is implemented only for greedy decoding as far as I remember. You could implement it yourself with modifying the logic in AbstractRNNTDecoding or extracting the alignments and manually building the time stamp list

So could I extract the alignments for beam search decoding and then building the timestamp list for each N full hypothesis (beam size = N)? Would this mean that with beam search decoding I could only extract N full hypotheses, or could I extract N words per word location and then build the time stamp list for each word location?

Yes that can be done. Currently it's not implemented in NeMo but you can calculate greedy word timestamps for each of the N hypothesis generated per sample. 

With the flags you have set, rnnt transcribe () will give you two items usually- first is the best hyp list and second is all hyp list (usually none unless you ser these flags). The second list is a list of NBest hypothesis which themselves contain a list of Hypothesize - each of which should have the alignment information. 

So with the following flags set:
decoder_type='rnnt'
compute_timestamps=True
rnnt_decoding.strategy='beam'
rnnt_decoding.beam.beam_size=4
rnnt_decoding.preserve_alignments=True
rnnt_decoding.beam.return_best_hypothesis=False

 the transcriptions object returned contains:

**1.** [0] = List of 2 elements, where: each element is a Hypothesis object, where: 
**1.1.** Element [0] has score=-4.7474, length=tensor(145) and has "timestep" and "alignments" fields; 
**1.2.** Element [1] has score=-10.1037, length=tensor(186) and also has "timestep" and "alignments" fields.
**2.** [1] = List of 2 elements, where: each element is a Hypothesis object, where: 
**2.1.** Element [0] is a list of 4 Hypothesis objects, where: 
**2.1.1.** Element [0] has score=-4.7474, length=tensor(145) and has "timestep" and "alignments" fields; 
**2.1.2.** Element [1] has score=-4.9974, length=tensor(145) and also has "timestep" and "alignments" fields; 
**2.1.3.** Element [2] has score=-4.9976, length=tensor(145) and also has "timestep" and "alignments" fields; 
**2.1.4.** Element [3] has score=-5.088, length=tensor(145) and also has "timestep" and "alignments" fields.
**2.2.** Element [1] is also a list of 4 Hypothesis objects, where:
**2.2.1.** Element [0] has score=-10.1037, length=tensor(186) and has "timestep" and "alignments" fields; 
**2.2.2.** Element [1] has score=-10.4386, length=tensor(186) and has "timestep" and "alignments" fields; 
**2.2.3.** Element [2] has score=-10.6843, length=tensor(186) and has "timestep" and "alignments" fields; 
**2.2.4.** Element [3] has score=-10.7804, length=tensor(186) and has "timestep" and "alignments" fields.

If I have a beam size of 4, and I want to time align each beam, what 4 Hypotheses with "timestep" and "alignments" fields do I need to use from the returned transcriptions object? I am not sure how the lists transcriptions[0] and transcriptions[1] differ from each other. Also, could you point me to where I can find the implementation of time aligning a Hypothesis object based on "timestep" and "alignments" fields, as I am not sure how to go about extracting word-level time-alignment information, where all I need is the start time and end time for each word in each beam hypothesis.
Also am I providing some redundant flags?

Ah sorry nevermind, I had a brief lapse of understanding. 

As I understand, transcriptions[0] contains a list of just the **best** Hypothesis per audio file, and transcriptions[1] contains a list per each audio file, where the lists contain **all** the Hypothesis objects returned by beam search decoding for that audio file. Also creating the timestamp information can be done by following https://colab.research.google.com/github/NVIDIA/NeMo/blob/stable/tutorials/asr/Offline_ASR.ipynb#scrollTo=_5hfsiDGsM19
All that's needed is to write a custom function to write the results to separate files per hypothesis. 

One question though: is the `rnnt_decoding.preserve_alignments=True` flag needed if I just want the timestamps for each word in each hypothesis? I don't see how it affects the generation of the `timestep` field of a Hypothesis object. Also, it seems like the `alignments` field of Hypothesis object is also not needed, what is its purpose? Or are the `timestep` fields created from the content of the `alignments` field?

Yes the preserve alignments field is needed. The rest of the timestamps and alignments are values derived if you set that flag. It won't work otherwise 

Thank you!