Repository navigation
Parakeet tdt - #44171
Parakeet tdt#44171
Conversation
Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline.
6c98cb8 to
f2b4938
Compare
There was a problem hiding this comment.
@lmaksym thank you putting together the PRs cleanly! I pushed a few changes for adapting to Transformers convention and added integration tests to compare with the original model from NeMo.
@hainan-xv and @nithinraok, your input could be useful for the TDT decoding, and also the loss computation.
- Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests
7f70c24 to
760b4b6
Compare
hainan-xv
left a comment
There was a problem hiding this comment.
Left a comment on the loss computation part.
ebezzam
left a comment
There was a problem hiding this comment.
@lmaksym thanks for porting the TDT loss! it's nice (1) to not have to depend on torchaudio and (2) to make the TDT loss available in Transformers!
It is functional with this example (single GPU): https://gist.github.com/ebezzam/6382bdabfc64bb2541ca9f77deb7678d#file-tdt_training_snippet-py
But quite slow...
I wonder if there is a custom gradient computation in NeMo? As I noticed in the paper (Section 3.1), they say "We derive an analytical solution for the gradient of the TDT loss, since automatic differentiation for transducer loss is highly inefficient."
FYI I can test/fix on my side for multi-GPU compatibility.
I'll look into that |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
run-slow: lasr, parakeet |
|
This comment contains models: ["models/lasr", "models/parakeet"] |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: auto, encodec, lasr, parakeet |
|
View the CircleCI Test Summary for this PR: https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=44171&sha=f97809 |
* parakeet tdt intergration * Add TDT decoder support for Parakeet ASR models Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline. * Add expected outputs for TDT, small fixes. * Separate CTC and TDT generate outputs. * Work with auto device, better init, * Test timestamps and expose token duration. * Add reproducer link. * fix: align TDT training and decoding with NeMo implementation - Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests * revert: restore lasr generated files to original state * warn: torchaudio rnnt_loss does not train duration head * Relax timestamp test, and test nits. * feat: TDT training * chore: for cuda detection and run without patching * Equivalent timestamp processing as Nemo, and various nits/cleanup. * Simplify durations config. * Update training examples. * chore: enable parralelism * chore: performance optimization * fix: formatting * Doc and testing nits * Use active mask from current step, and nits. * Better pre-allocate. * TDT has separate pad token and blank token. * Regenerate lasr. * Style checks and nits * Nits, put back ctc loss test * More standard model output. * Style * Remove compute_loss flag and allow monkey patching to tdt loss * Update src/transformers/models/parakeet/modular_parakeet.py Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> * Address various comments. * More compatible with Transformers forward/generate approach * compile option for generation and decoder cache * Cleaner, better conventions. * Update with main. * doc nits * Imitate whisper for encoder outputs as input * Address tests and nits. * Inherit from GenerateMixIn for get_compiled_call * Comment nit * forward cleanup * generate cleanup + separate generation file * generate: add _supported_generation_modes * automatic init of the loss * modular cleanups * use is_encoder_decoder * timestamp processing fully from tokens + durations * convertion script update * test update * make * test update * test update * ensure correct loss computation * kernel loss * test loss integration * push to hub pr * integration tests to rely fully on transcripts * udpate fixtures * we don't need to monkey patch with numba anymore! * fix pipeline usage * nit * fix usage * Pass through tests and examples: improve kernel fallback, update with nvidia checkpoint, style checks. * Update checkpoint * Add TDT to mapping after merge. * Fix lasr generate test. * Output attention mask if labels provided for computing loss. * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Improve ParakeetTDTDecoderCache definition and usage. * Remove tuple parsing. * processor refactor * Update conversion. * Remove kernel to address in separate PR, modular nit. * Lasr modular * Add attention mask docstring. --------- Co-authored-by: Hainan Xu <hainanx@nvidia.com> Co-authored-by: Eric B <ebezzam@gmail.com> Co-authored-by: Eric Bezzam <4757445+ebezzam@users.noreply.github.com> Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Tarek Ziade <tarek@ziade.org>
* parakeet tdt intergration * Add TDT decoder support for Parakeet ASR models Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline. * Add expected outputs for TDT, small fixes. * Separate CTC and TDT generate outputs. * Work with auto device, better init, * Test timestamps and expose token duration. * Add reproducer link. * fix: align TDT training and decoding with NeMo implementation - Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests * revert: restore lasr generated files to original state * warn: torchaudio rnnt_loss does not train duration head * Relax timestamp test, and test nits. * feat: TDT training * chore: for cuda detection and run without patching * Equivalent timestamp processing as Nemo, and various nits/cleanup. * Simplify durations config. * Update training examples. * chore: enable parralelism * chore: performance optimization * fix: formatting * Doc and testing nits * Use active mask from current step, and nits. * Better pre-allocate. * TDT has separate pad token and blank token. * Regenerate lasr. * Style checks and nits * Nits, put back ctc loss test * More standard model output. * Style * Remove compute_loss flag and allow monkey patching to tdt loss * Update src/transformers/models/parakeet/modular_parakeet.py Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> * Address various comments. * More compatible with Transformers forward/generate approach * compile option for generation and decoder cache * Cleaner, better conventions. * Update with main. * doc nits * Imitate whisper for encoder outputs as input * Address tests and nits. * Inherit from GenerateMixIn for get_compiled_call * Comment nit * forward cleanup * generate cleanup + separate generation file * generate: add _supported_generation_modes * automatic init of the loss * modular cleanups * use is_encoder_decoder * timestamp processing fully from tokens + durations * convertion script update * test update * make * test update * test update * ensure correct loss computation * kernel loss * test loss integration * push to hub pr * integration tests to rely fully on transcripts * udpate fixtures * we don't need to monkey patch with numba anymore! * fix pipeline usage * nit * fix usage * Pass through tests and examples: improve kernel fallback, update with nvidia checkpoint, style checks. * Update checkpoint * Add TDT to mapping after merge. * Fix lasr generate test. * Output attention mask if labels provided for computing loss. * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Improve ParakeetTDTDecoderCache definition and usage. * Remove tuple parsing. * processor refactor * Update conversion. * Remove kernel to address in separate PR, modular nit. * Lasr modular * Add attention mask docstring. --------- Co-authored-by: Hainan Xu <hainanx@nvidia.com> Co-authored-by: Eric B <ebezzam@gmail.com> Co-authored-by: Eric Bezzam <4757445+ebezzam@users.noreply.github.com> Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Tarek Ziade <tarek@ziade.org>
* parakeet tdt intergration * Add TDT decoder support for Parakeet ASR models Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline. * Add expected outputs for TDT, small fixes. * Separate CTC and TDT generate outputs. * Work with auto device, better init, * Test timestamps and expose token duration. * Add reproducer link. * fix: align TDT training and decoding with NeMo implementation - Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests * revert: restore lasr generated files to original state * warn: torchaudio rnnt_loss does not train duration head * Relax timestamp test, and test nits. * feat: TDT training * chore: for cuda detection and run without patching * Equivalent timestamp processing as Nemo, and various nits/cleanup. * Simplify durations config. * Update training examples. * chore: enable parralelism * chore: performance optimization * fix: formatting * Doc and testing nits * Use active mask from current step, and nits. * Better pre-allocate. * TDT has separate pad token and blank token. * Regenerate lasr. * Style checks and nits * Nits, put back ctc loss test * More standard model output. * Style * Remove compute_loss flag and allow monkey patching to tdt loss * Update src/transformers/models/parakeet/modular_parakeet.py Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> * Address various comments. * More compatible with Transformers forward/generate approach * compile option for generation and decoder cache * Cleaner, better conventions. * Update with main. * doc nits * Imitate whisper for encoder outputs as input * Address tests and nits. * Inherit from GenerateMixIn for get_compiled_call * Comment nit * forward cleanup * generate cleanup + separate generation file * generate: add _supported_generation_modes * automatic init of the loss * modular cleanups * use is_encoder_decoder * timestamp processing fully from tokens + durations * convertion script update * test update * make * test update * test update * ensure correct loss computation * kernel loss * test loss integration * push to hub pr * integration tests to rely fully on transcripts * udpate fixtures * we don't need to monkey patch with numba anymore! * fix pipeline usage * nit * fix usage * Pass through tests and examples: improve kernel fallback, update with nvidia checkpoint, style checks. * Update checkpoint * Add TDT to mapping after merge. * Fix lasr generate test. * Output attention mask if labels provided for computing loss. * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Improve ParakeetTDTDecoderCache definition and usage. * Remove tuple parsing. * processor refactor * Update conversion. * Remove kernel to address in separate PR, modular nit. * Lasr modular * Add attention mask docstring. --------- Co-authored-by: Hainan Xu <hainanx@nvidia.com> Co-authored-by: Eric B <ebezzam@gmail.com> Co-authored-by: Eric Bezzam <4757445+ebezzam@users.noreply.github.com> Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Tarek Ziade <tarek@ziade.org>
* parakeet tdt intergration * Add TDT decoder support for Parakeet ASR models Implement Token-and-Duration Transducer (TDT) decoding for Parakeet models, extending the existing CTC-only support. This adds ParakeetForTDT with greedy TDT decoding in generate(), per-token timestamp generation, and full integration with AutoModelForTDT, processors, and ASR pipeline. * Add expected outputs for TDT, small fixes. * Separate CTC and TDT generate outputs. * Work with auto device, better init, * Test timestamps and expose token duration. * Add reproducer link. * fix: align TDT training and decoding with NeMo implementation - Use -100 label padding for training (HF convention) - Fix timestamp recording in inner blank-seeking loop - Add max_symbols_per_step guard matching NeMo - Clean up decoding loop - Add TDT training example to docs - Use setUpClass for TDT integration tests * revert: restore lasr generated files to original state * warn: torchaudio rnnt_loss does not train duration head * Relax timestamp test, and test nits. * feat: TDT training * chore: for cuda detection and run without patching * Equivalent timestamp processing as Nemo, and various nits/cleanup. * Simplify durations config. * Update training examples. * chore: enable parralelism * chore: performance optimization * fix: formatting * Doc and testing nits * Use active mask from current step, and nits. * Better pre-allocate. * TDT has separate pad token and blank token. * Regenerate lasr. * Style checks and nits * Nits, put back ctc loss test * More standard model output. * Style * Remove compute_loss flag and allow monkey patching to tdt loss * Update src/transformers/models/parakeet/modular_parakeet.py Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> * Address various comments. * More compatible with Transformers forward/generate approach * compile option for generation and decoder cache * Cleaner, better conventions. * Update with main. * doc nits * Imitate whisper for encoder outputs as input * Address tests and nits. * Inherit from GenerateMixIn for get_compiled_call * Comment nit * forward cleanup * generate cleanup + separate generation file * generate: add _supported_generation_modes * automatic init of the loss * modular cleanups * use is_encoder_decoder * timestamp processing fully from tokens + durations * convertion script update * test update * make * test update * test update * ensure correct loss computation * kernel loss * test loss integration * push to hub pr * integration tests to rely fully on transcripts * udpate fixtures * we don't need to monkey patch with numba anymore! * fix pipeline usage * nit * fix usage * Pass through tests and examples: improve kernel fallback, update with nvidia checkpoint, style checks. * Update checkpoint * Add TDT to mapping after merge. * Fix lasr generate test. * Output attention mask if labels provided for computing loss. * Apply suggestion from @ArthurZucker Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> * Improve ParakeetTDTDecoderCache definition and usage. * Remove tuple parsing. * processor refactor * Update conversion. * Remove kernel to address in separate PR, modular nit. * Lasr modular * Add attention mask docstring. --------- Co-authored-by: Hainan Xu <hainanx@nvidia.com> Co-authored-by: Eric B <ebezzam@gmail.com> Co-authored-by: Eric Bezzam <4757445+ebezzam@users.noreply.github.com> Co-authored-by: eustlb <94853470+eustlb@users.noreply.github.com> Co-authored-by: Arthur <48595927+ArthurZucker@users.noreply.github.com> Co-authored-by: Tarek Ziade <tarek@ziade.org>
## What does this PR do? Goal of this PR is to add a [kernel](https://huggingface.co/eustlb/tdt-loss/tree/v1) that has been drafted by @eustlb for the recently merged [Parakeet TDT](huggingface/transformers#44171). It's a very rough draft with a lot of Claude help as it's my first time working with kernels 😅 TODO - [ ] @eustlb: I guess you had a .cu file to create the .so file [here](https://huggingface.co/eustlb/tdt-loss/tree/v1/build/torch211-cxx11-cu128-x86_64-linux/tdt_loss)? From what I understand, we will need to add it here. - [ ] Add the kernel to https://huggingface.co/kernels-community - [ ] Update the tests - [ ] Double-check with [this checklist](https://github.com/huggingface/kernels-community/blob/main/CONTRIBUTING.md#how-to-add-a-new-kernel) Perhaps we can see how other kernels [used in Transformers](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/hub_kernels.py) were integrated. As the kernel is for a function, maybe [rotary_pos_emb](https://github.com/huggingface/transformers/blob/0df9b7fcaab447c75543598e6d959065c2296a24/src/transformers/integrations/hub_kernels.py#L221) is a good example? [Here](https://github.com/huggingface/kernels-community/tree/main/rotary) are the files in this repo. Here is a draft PR for adding/using this kernel in Transformers: huggingface/transformers#46048 cc @vasqu --------- Co-authored-by: Deep-unlearning <steven@huggingface.co> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
What does this PR do?
This PR adds TDT decoder support for Parakeet ASR models, extending the existing CTC-only implementation.
It incorporates the initial TDT integration work from #41545 by @hainan-xv (was not merged) and and addresses all review feedback from both #41545 and #43357.
Changes
ParakeetForTDTmodel with greedy TDT decoding ingenerate()ParakeetTDTDecoder(LSTM prediction network) andParakeetTDTJointNetworkasnn.Modulesubclassesreturn_timestamps=TrueAutoModelForTDTauto class with pipeline, processor, and tokenizer integrationParakeetTDTConfigmatching the CTC pattern (no nested decoder/joint configs)ParakeetPreTrainedModelbase between CTC and TDT (no separate TDT base class)Validation
make check-repopassesBefore submitting
Pull Request section?
documentation guidelines, and
here are tips on formatting docstrings.
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@ebezzam and @hainan-xv please review
-->