TDT confidence fix

[GNroy]

What does this PR do ?

This PR enables confidence scores computing for TDT models.

It introduces a new ConfidenceConfig parameter named tdt_include_duration which allows caclulating frame-level confidence for the token duration prediction.

Collection: [ASR]

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Jenkins CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

There's no need to comment jenkins on the PR to trigger Jenkins CI.
The GitHub Actions CI will run automatically when the PR is opened.
To run CI on an untrusted fork, a NeMo user with write access must click "Approve and run".

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 parakeet-tdt-1.1b unable to calculate ASR confidence scores! #8737

[artbataev]

Please, add comments what's going on here and why

[GNroy]

There is a comment a few lines above # there is only one score per frame most of the time.
Apparently, it's not enough. The idea is to skip running for loop for frames with a single prediction.

[artbataev]

I found 2 problems with the current implementation:

1. when decoding with duration confidence, inside GreedyBatchedTDTLoopLabelsComputer include_duration_confidence parameter is False (you can print it)

python examples/asr/speech_to_text_eval.py \
 pretrained_name=nvidia/parakeet-tdt-1.1b \
 dataset_manifest=~/code/asr_particles/manifests/librispeech/test_other.json \
 batch_size=16 \
 output_filename=test_other_decoded.jsonl \
 amp=false \
 amp_dtype=bfloat16 \
 rnnt_decoding.greedy.tdt_include_duration_confidence=true \
 rnnt_decoding.greedy.preserve_frame_confidence=true \
 preserve_alignment=true

• I think that also you need to adjust BatchedAlignments._allocate_more method to correctly reflect the changed shape self.frame_confidence = torch.cat((self.frame_confidence, torch.zeros_like(self.frame_confidence)), dim=1) (should be dim=1)

2. the model fails with bf16 (amp=true)

python examples/asr/speech_to_text_eval.py \
 pretrained_name=nvidia/parakeet-tdt-1.1b \
 dataset_manifest=~/code/asr_particles/manifests/librispeech/test_other.json \
 batch_size=16 \
 output_filename=test_other_decoded.jsonl \
 amp=true \
 amp_dtype=bfloat16 \
 rnnt_decoding.greedy.tdt_include_duration_confidence=true \
 rnnt_decoding.greedy.preserve_frame_confidence=true \
 preserve_alignment=true

[artbataev]

Use torch.stack here and below to concat along new dimension (instead of multiple unsqueeze)

[GNroy]

Thanks, done.

[GNroy]

I found 2 problems with the current implementation:

I think I fixed those problems. Could you please check again?

[artbataev]

LGTM, thanks!