Streaming decoding policies (Wait-K and AlignAtt) for Canary model

[andrusenkoau]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Add streaming decoding policies support (Wait-K and AlignAtt) for Canary model.

Wait-K policy predicts only one token per each new speech chunk that increases the overall latency - original paper. In this case, it is unclear at what point you can forget part of the left context when recognizing with a limited window. It is recommended to set the left context to the maximum possible value (infinite left context) for the waitk policy.

AlignAtt policy predicts tokens according to the cross-attention condition - original paper. If the condition is met, then the audio size does not need to be increased, and the prediction of the next token continues. Otherwise, the audio buffer size needs to be increased. This policy shows lower latency in comparison with waitk. This policy is also suitable for window recognition (left context is fixed and not infinite), but you can lose accuracy in this case.

Collection: [ASR]

Changelog

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

Usage

Decoding policy is controlled by AEDStreamingDecodingConfig. You can choose waitk or alignatt policy. The remaining parameters (such as alignatt_thr or waitk_lagging) need to be selected depending on the data and the task (for example, for AST, you can increase waitk_lagging, which works for both policies).

Remember to manage prompt parameters using +prompt (for example, +prompt.pnc=yes/no, +prompt.task=asr/ast, +prompt.source_lang=en, +prompt.target_lang=de, and so on). This is especially important for AST task.

python3 ${NEMO_DIR}/examples/asr/asr_chunked_inference/aed/speech_to_text_aed_streaming_infer.py \
    pretrained_name=${ASR_MODEL_NAME} \
    dataset_manifest=${DECODING_MANIFEST} \
    output_filename=${OUTPUT_MANIFEST} \
    left_context_secs=10 \
    chunk_secs=1 \
    right_context_secs=0.5 \
    batch_size=32\
    decoding.streaming_policy=waitk \ # [waitk or alignatt]
    decoding.alignatt_thr=8 \
    decoding.waitk_lagging=2 \
    decoding.exclude_sink_frames=8 \
    decoding.xatt_scores_layer=-2 \
    decoding.hallucinations_detector=True \
    +prompt.pnc=yes \
    +prompt.task=asr

GitHub Actions CI

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

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first 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 # (issue)

[artbataev]

Thanks a lot! Great work!

[nithinraok]

Thanks @andrusenkoau LGTM.

Excellent work!