diff --git a/egs/europarl_st/SRT/LICENSE b/egs/europarl_st/SRT/LICENSE
new file mode 100644
index 0000000000..0a92d40a74
--- /dev/null
+++ b/egs/europarl_st/SRT/LICENSE
@@ -0,0 +1,201 @@
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to the Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by the Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding any notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. Please also get
+ "http://www.apache.org/licenses/" on a line by itself so that
+ the browser will use it as a clickable link.
+
+ Copyright 2026 Nanjie Li
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
diff --git a/egs/europarl_st/SRT/README.md b/egs/europarl_st/SRT/README.md
new file mode 100644
index 0000000000..5c32e021fc
--- /dev/null
+++ b/egs/europarl_st/SRT/README.md
@@ -0,0 +1,179 @@
+# LCMA-SRT: Language-Conditional Mixture-of-Experts Adapters for Joint Multilingual Speech Recognition and Translation
+
+[](LICENSE)
+[](https://aclanthology.org/2026.acl-long.1634/)
+
+Neural transducers offer an alignment-free framework for speech-to-text modeling, and hierarchical transducer architectures further improve multilingual joint automatic speech recognition (ASR) and speech translation (ST) by stacking a translation-focused encoder on top of an ASR encoder. However, extending hierarchical transducers to multilingual many-to-many settings remains challenging: fully shared models often suffer from negative transfer and unstable target-language generation, while training separate models for each direction is computationally prohibitive. We propose LCMA-SRT (Language-Conditional Mixture-of-Experts Adapters for Speech Recognition and Translation), which augments a hierarchical transducer with language-conditional Mixture-of-Experts (MoE) adapters. A source-conditioned MoE adapter (SRC-MoE) uses source-language embeddings to reduce cross-language interference and improve multilingual ASR. A target-conditioned MoE adapter (TGT-MoE) uses the desired target language to reduce cross-target interference and stabilize target-language generation in many-to-many ST. Experiments on Europarl-ST (9 languages, 72 directions) show that LCMA-SRT improves both ASR and ST within a single joint model, reducing average WER and improving BLEU and COMET over strong hierarchical transducer baselines.
+
+
+
+## Installation
+
+Please refer to the [icefall installation guide](https://k2-fsa.github.io/icefall/installation/index.html).
+
+## Data Preparation
+
+Download the Europarl-ST dataset from the [official page](https://www.mllp.upv.es/europarl-st/).
+
+See [local](local) for data preparation scripts.
+
+## Training
+
+### Stage 1: Multilingual ASR Pretraining
+
+```bash
+bash lcma_srt/train/stage1/cr_ctc_sc_moe.sh
+```
+
+### Stage 2: Many-to-Many Joint Training
+
+```bash
+bash lcma_srt/train/stage2/lcma_srt.sh
+```
+
+## Decoding
+
+### Stage 1: ASR Decoding
+
+```bash
+bash lcma_srt/decode/stage1/decode_cr_ctc_sc_moe.sh
+```
+
+### Stage 2: Joint ASR+ST Decoding
+
+```bash
+bash lcma_srt/decode/stage2/decode_lcma_srt.sh
+```
+
+## Checkpoint
+
+Pre-trained model checkpoints are available on [Hugging Face](https://huggingface.co/linanjie0820/lcma-srt).
+
+## Main Results
+
+### Multilingual ASR Pretraining
+
+
+
+
+ | Model |
+ WER (%) ↓ |
+
+
+ | de | en | es | fr | it |
+ nl | pl | pt | ro | Avg |
+
+
+
+
+ | CR-CTC |
+ 24.57 | 18.59 | 20.76 | 19.24 | 17.33 |
+ 36.75 | 25.28 | 19.82 | 18.77 | 22.35 |
+
+
+ | + MoE |
+ 24.39 | 18.41 | 20.16 | 18.61 | 17.28 |
+ 36.83 | 24.36 | 19.70 | 18.79 | 22.06 |
+
+
+ | + S-Bias |
+ 23.89 | 17.60 | 19.58 | 17.41 | 16.73 |
+ 34.72 | 23.63 | 18.21 | 17.97 | 21.08 |
+
+
+ | + SRC-MoE |
+ 23.34 | 17.45 | 19.41 |
+ 17.34 | 16.27 | 35.20 |
+ 23.28 | 18.16 | 17.48 |
+ 20.88 |
+
+
+
+
+### Many-to-Many Joint Training (Average)
+
+
+
+
+ | Model |
+ WER (%)↓ |
+ Average BLEU ↑ |
+
+
+ | de | en | es | fr | it | nl | pl | pt | ro | Avg |
+
+
+
+
+ | HENT-SRT-M20×9 |
+ 23.28 |
+ 10.7 | 21.2 | 19.1 | 18.2 | 14.2 | 16.5 | 7.2 | 18.4 | 12.1 | 15.3 |
+
+
+ | HENT-SRT-M2M |
+ 16.65 |
+ 2.6 | 12.8 | 5.5 | 4.0 | 1.8 | 3.5 | 1.2 | 4.9 | 2.5 | 4.3 |
+
+
+ | LCMA-SRT |
+ 15.71 |
+ 15.2 | 25.9 | 25.8 | 24.7 |
+ 20.0 | 20.5 | 10.7 | 23.9 | 17.6 | 20.5 |
+
+
+
+
+
+
+
+ | Model |
+ LMR (%)↓ |
+ Average COMET ↑ |
+
+
+ | de | en | es | fr | it | nl | pl | pt | ro | Avg |
+
+
+
+
+ | HENT-SRT-M20×9 |
+ 0.65 |
+ 0.507 | 0.656 | 0.587 | 0.542 | 0.565 | 0.558 | 0.550 | 0.609 | 0.598 | 0.575 |
+
+
+ | HENT-SRT-M2M |
+ 84.95 |
+ 0.380 | 0.543 | 0.478 | 0.427 | 0.435 | 0.401 | 0.385 | 0.471 | 0.406 | 0.436 |
+
+
+ | LCMA-SRT |
+ 0.75 |
+ 0.574 | 0.715 | 0.682 | 0.627 |
+ 0.656 | 0.613 | 0.616 | 0.693 | 0.678 | 0.651 |
+
+
+
+
+For complete per-direction results (all 72 directions), see [RESULTS.md](RESULTS.md).
+
+## Evaluation
+
+We use BLEU for surface-level matching, COMET for semantic adequacy, and sentence-level target-language mismatch rate (LMR) using an off-the-shelf language identification [model](https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin), where a hypothesis is counted as matched only if it is classified as the specified target language with confidence >= 0.7. ASR performance is assessed using word error rate (WER).
+
+## Citation
+
+If this project is useful for your research, please cite:
+
+```bibtex
+@inproceedings{li2026lcma,
+ title={LCMA-SRT: Language-Conditional Mixture-of-Experts Adapters for Joint Multilingual Speech Recognition and Translation},
+ author={Li, Nanjie and Guo, Xiaoyong and Huang, Hao and Haihua, Xu and Shi, Wei},
+ booktitle={Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
+ pages={35363--35377},
+ year={2026}
+}
+```
+
+## License
+
+This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details.
diff --git a/egs/europarl_st/SRT/RESULTS.md b/egs/europarl_st/SRT/RESULTS.md
new file mode 100644
index 0000000000..1c13f9ab9f
--- /dev/null
+++ b/egs/europarl_st/SRT/RESULTS.md
@@ -0,0 +1,918 @@
+# Detailed Results
+
+This file contains the complete per-direction results for all 72 translation directions on Europarl-ST.
+
+## Many-to-Many Joint Training (All Directions)
+
+### WER
+
+
+
+ | SRC\TGT |
+ Model |
+ WER (%) ↓ |
+
+
+ | de | en | es | fr | it | nl | pl | pt | ro |
+
+
+
+
+
+ | de |
+ HENT-SRT-M20×9 |
+ - | 21.80 | 26.64 | 27.12 | 27.44 | 26.43 | 26.51 | 26.51 | 26.62 |
+
+
+ | HENT-SRT-M2M |
+ - | 19.09 | 18.77 | 18.82 | 19.09 | 18.86 | 18.79 | 18.99 | 18.86 |
+
+
+ | LCMA-SRT |
+ - | 18.01 | 17.84 | 17.85 | 18.23 | 17.92 | 17.75 | 17.93 | 17.99 |
+
+
+ | TGT-MoE→MoE |
+ - | 18.83 | 18.76 | 18.75 | 18.92 | 18.74 | 18.61 | 18.81 | 18.85 |
+
+
+ | TGT-MoE→T-Bias |
+ - | 18.13 | 17.88 | 17.86 | 18.16 | 17.92 | 17.80 | 17.93 | 17.97 |
+
+
+ | w/o TGT-MoE |
+ - | 18.83 | 18.59 | 18.62 | 18.95 | 18.60 | 18.44 | 18.57 | 18.78 |
+
+
+ | w/o SRC-MoE |
+ - | 18.66 | 18.35 | 18.33 | 18.65 | 18.41 | 18.30 | 18.51 | 18.50 |
+
+
+
+
+
+ | en |
+ HENT-SRT-M20×9 |
+ 16.42 | - | 17.32 | 17.56 | 17.08 | 17.45 | 17.29 | 17.18 | 17.21 |
+
+
+ | HENT-SRT-M2M |
+ 13.92 | - | 13.94 | 14.02 | 13.79 | 13.84 | 13.99 | 13.90 | 13.53 |
+
+
+ | LCMA-SRT |
+ 12.94 | - | 12.93 | 13.02 | 12.84 | 12.87 | 12.92 | 12.97 | 12.64 |
+
+
+ | TGT-MoE→MoE |
+ 13.28 | - | 13.27 | 13.33 | 13.14 | 13.18 | 13.31 | 13.34 | 12.88 |
+
+
+ | TGT-MoE→T-Bias |
+ 13.26 | - | 13.25 | 13.35 | 13.05 | 13.16 | 13.26 | 13.25 | 12.94 |
+
+
+ | w/o TGT-MoE |
+ 13.41 | - | 13.38 | 13.45 | 13.26 | 13.31 | 13.43 | 13.39 | 13.04 |
+
+
+ | w/o SRC-MoE |
+ 13.50 | - | 13.49 | 13.61 | 13.46 | 13.39 | 13.53 | 13.48 | 13.13 |
+
+
+
+
+
+ | es |
+ HENT-SRT-M20×9 |
+ 21.29 | 17.77 | - | 22.25 | 22.83 | 22.68 | 21.93 | 22.82 | 22.66 |
+
+
+ | HENT-SRT-M2M |
+ 15.96 | 15.80 | - | 15.91 | 15.69 | 15.97 | 15.85 | 15.89 | 15.75 |
+
+
+ | LCMA-SRT |
+ 15.30 | 15.14 | - | 15.27 | 15.02 | 15.31 | 15.26 | 15.25 | 15.14 |
+
+
+ | TGT-MoE→MoE |
+ 15.61 | 15.51 | - | 15.65 | 15.41 | 15.60 | 15.56 | 15.52 | 15.57 |
+
+
+ | TGT-MoE→T-Bias |
+ 15.19 | 15.13 | - | 15.21 | 15.03 | 15.21 | 15.14 | 15.11 | 14.89 |
+
+
+ | w/o TGT-MoE |
+ 15.96 | 15.81 | - | 15.97 | 15.74 | 15.99 | 15.89 | 15.90 | 15.82 |
+
+
+ | w/o SRC-MoE |
+ 15.46 | 15.22 | - | 15.41 | 15.16 | 15.46 | 15.34 | 15.35 | 15.11 |
+
+
+
+
+
+ | fr |
+ HENT-SRT-M20×9 |
+ 19.37 | 16.07 | 19.82 | - | 20.41 | 19.30 | 19.45 | 19.86 | 20.80 |
+
+
+ | HENT-SRT-M2M |
+ 13.40 | 13.38 | 13.28 | - | 13.42 | 13.36 | 13.39 | 13.38 | 13.37 |
+
+
+ | LCMA-SRT |
+ 12.58 | 12.51 | 12.53 | - | 12.51 | 12.50 | 12.56 | 12.55 | 12.65 |
+
+
+ | TGT-MoE→MoE |
+ 13.15 | 13.10 | 13.02 | - | 13.08 | 12.99 | 13.15 | 13.16 | 13.27 |
+
+
+ | TGT-MoE→T-Bias |
+ 12.53 | 12.49 | 12.37 | - | 12.47 | 12.45 | 12.48 | 12.52 | 12.51 |
+
+
+ | w/o TGT-MoE |
+ 12.75 | 12.77 | 12.62 | - | 12.70 | 12.67 | 12.69 | 12.65 | 12.78 |
+
+
+ | w/o SRC-MoE |
+ 12.58 | 12.64 | 12.49 | - | 12.55 | 12.65 | 12.56 | 12.64 | 12.74 |
+
+
+
+
+
+ | it |
+ HENT-SRT-M20×9 |
+ 18.18 | 15.05 | 19.19 | 19.32 | - | 19.06 | 18.60 | 19.00 | 19.91 |
+
+
+ | HENT-SRT-M2M |
+ 13.10 | 13.19 | 13.13 | 13.24 | - | 13.17 | 12.98 | 13.18 | 13.27 |
+
+
+ | LCMA-SRT |
+ 12.50 | 12.41 | 12.52 | 12.63 | - | 12.59 | 12.42 | 12.62 | 12.66 |
+
+
+ | TGT-MoE→MoE |
+ 13.00 | 12.92 | 13.03 | 13.04 | - | 13.05 | 12.89 | 13.11 | 13.27 |
+
+
+ | TGT-MoE→T-Bias |
+ 12.67 | 12.63 | 12.77 | 12.80 | - | 12.74 | 12.67 | 12.84 | 12.95 |
+
+
+ | w/o TGT-MoE |
+ 12.91 | 12.91 | 12.97 | 13.03 | - | 12.97 | 12.85 | 13.13 | 13.08 |
+
+
+ | w/o SRC-MoE |
+ 12.92 | 12.90 | 12.96 | 13.07 | - | 13.07 | 12.84 | 13.12 | 13.13 |
+
+
+
+
+
+ | nl |
+ HENT-SRT-M20×9 |
+ 38.99 | 32.95 | 38.85 | 38.85 | 39.52 | - | 38.99 | 39.32 | 39.26 |
+
+
+ | HENT-SRT-M2M |
+ 28.59 | 28.65 | 28.73 | 28.46 | 28.62 | - | 28.46 | 28.46 | 28.47 |
+
+
+ | LCMA-SRT |
+ 27.01 | 27.23 | 26.89 | 26.91 | 27.20 | - | 26.93 | 27.07 | 26.82 |
+
+
+ | TGT-MoE→MoE |
+ 28.47 | 28.60 | 28.57 | 28.38 | 28.58 | - | 28.34 | 28.47 | 28.34 |
+
+
+ | TGT-MoE→T-Bias |
+ 27.33 | 27.39 | 27.28 | 27.17 | 27.57 | - | 27.32 | 27.29 | 27.26 |
+
+
+ | w/o TGT-MoE |
+ 28.71 | 28.80 | 28.57 | 28.56 | 28.75 | - | 28.61 | 28.52 | 28.48 |
+
+
+ | w/o SRC-MoE |
+ 27.85 | 28.02 | 27.74 | 27.69 | 27.94 | - | 27.65 | 27.77 | 27.55 |
+
+
+
+
+
+ | pl |
+ HENT-SRT-M20×9 |
+ 25.89 | 22.01 | 26.33 | 27.19 | 25.99 | 26.47 | - | 27.13 | 27.36 |
+
+
+ | HENT-SRT-M2M |
+ 18.26 | 18.27 | 18.14 | 18.21 | 17.87 | 18.27 | - | 18.29 | 18.00 |
+
+
+ | LCMA-SRT |
+ 17.54 | 17.39 | 17.32 | 17.36 | 17.01 | 17.43 | - | 17.57 | 17.11 |
+
+
+ | TGT-MoE→MoE |
+ 18.10 | 17.96 | 17.97 | 18.07 | 17.47 | 18.01 | - | 18.14 | 17.67 |
+
+
+ | TGT-MoE→T-Bias |
+ 17.56 | 17.32 | 17.42 | 17.45 | 17.06 | 17.41 | - | 17.50 | 17.37 |
+
+
+ | w/o TGT-MoE |
+ 18.30 | 18.14 | 18.02 | 18.18 | 17.85 | 18.17 | - | 18.24 | 17.92 |
+
+
+ | w/o SRC-MoE |
+ 17.88 | 17.55 | 17.76 | 17.79 | 17.57 | 17.79 | - | 18.00 | 17.70 |
+
+
+
+
+
+ | pt |
+ HENT-SRT-M20×9 |
+ 19.90 | 16.27 | 21.74 | 20.82 | 20.77 | 20.99 | 20.48 | - | 20.53 |
+
+
+ | HENT-SRT-M2M |
+ 13.60 | 13.59 | 13.52 | 13.59 | 13.38 | 13.58 | 13.57 | - | 13.34 |
+
+
+ | LCMA-SRT |
+ 12.37 | 12.72 | 12.28 | 12.37 | 12.08 | 12.38 | 12.40 | - | 12.19 |
+
+
+ | TGT-MoE→MoE |
+ 13.15 | 13.39 | 13.10 | 13.18 | 12.96 | 13.16 | 13.12 | - | 12.91 |
+
+
+ | TGT-MoE→T-Bias |
+ 12.53 | 12.72 | 12.50 | 12.55 | 12.30 | 12.52 | 12.57 | - | 12.30 |
+
+
+ | w/o TGT-MoE |
+ 13.12 | 13.29 | 13.01 | 13.12 | 12.95 | 13.09 | 13.06 | - | 12.82 |
+
+
+ | w/o SRC-MoE |
+ 12.75 | 12.86 | 12.63 | 12.75 | 12.49 | 12.77 | 12.67 | - | 12.48 |
+
+
+
+
+
+ | ro |
+ HENT-SRT-M20×9 |
+ 22.32 | 15.85 | 21.87 | 22.04 | 23.97 | 22.88 | 23.63 | 22.82 | - |
+
+
+ | HENT-SRT-M2M |
+ 14.59 | 14.20 | 14.52 | 14.42 | 14.17 | 14.59 | 14.49 | 14.65 | - |
+
+
+ | LCMA-SRT |
+ 13.64 | 13.29 | 13.51 | 13.46 | 13.38 | 13.61 | 13.54 | 13.72 | - |
+
+
+ | TGT-MoE→MoE |
+ 14.64 | 14.34 | 14.54 | 14.40 | 14.36 | 14.62 | 14.56 | 14.67 | - |
+
+
+ | TGT-MoE→T-Bias |
+ 13.99 | 13.63 | 13.93 | 13.82 | 13.62 | 13.98 | 13.88 | 14.01 | - |
+
+
+ | w/o TGT-MoE |
+ 15.04 | 14.96 | 14.91 | 14.92 | 14.77 | 15.05 | 14.99 | 15.11 | - |
+
+
+ | w/o SRC-MoE |
+ 14.09 | 13.74 | 13.97 | 13.88 | 13.73 | 14.05 | 13.91 | 14.13 | - |
+
+
+
+
+
+### BLEU
+
+
+
+
+ | SRC\TGT |
+ Model |
+ BLEU ↑ |
+
+
+ | de | en | es | fr | it | nl | pl | pt | ro |
+
+
+
+
+
+ | de |
+ HENT-SRT-M20×9 |
+ - | 17.5 | 13.3 | 12.1 | 8.7 | 16.2 | 5.9 | 12.4 | 8.3 |
+
+
+ | HENT-SRT-M2M |
+ - | 11.0 | 3.7 | 3.3 | 1.1 | 4.1 | 1.6 | 4.0 | 2.2 |
+
+
+ | LCMA-SRT |
+ - | 22.0 | 19.7 | 20.2 | 14.5 | 19.0 | 8.9 | 18.7 | 13.5 |
+
+
+ | TGT-MoE→MoE |
+ - | 12.4 | 3.2 | 2.1 | 1.0 | 2.8 | 1.4 | 3.7 | 2.1 |
+
+
+ | TGT-MoE→T-Bias |
+ - | 19.6 | 18.0 | 18.0 | 12.5 | 16.5 | 6.9 | 17.0 | 11.5 |
+
+
+ | w/o TGT-MoE |
+ - | 11.4 | 4.1 | 3.0 | 1.0 | 3.4 | 1.5 | 4.0 | 2.2 |
+
+
+ | w/o SRC-MoE |
+ - | 21.4 | 19.2 | 19.8 | 13.9 | 18.9 | 8.7 | 18.5 | 13.5 |
+
+
+
+
+
+ | en |
+ HENT-SRT-M20×9 |
+ 15.4 | - | 26.0 | 24.6 | 19.0 | 21.9 | 9.7 | 23.1 | 19.8 |
+
+
+ | HENT-SRT-M2M |
+ 4.0 | - | 9.7 | 6.5 | 3.1 | 5.3 | 1.6 | 7.1 | 4.5 |
+
+
+ | LCMA-SRT |
+ 20.1 | - | 33.4 | 30.7 | 25.0 | 25.4 | 14.7 | 29.4 | 26.3 |
+
+
+ | TGT-MoE→MoE |
+ 3.8 | - | 9.9 | 6.0 | 3.1 | 4.6 | 1.3 | 7.5 | 4.1 |
+
+
+ | TGT-MoE→T-Bias |
+ 17.4 | - | 30.0 | 27.3 | 22.4 | 22.1 | 11.2 | 26.2 | 21.1 |
+
+
+ | w/o TGT-MoE |
+ 3.1 | - | 10.6 | 6.1 | 2.9 | 4.5 | 1.4 | 7.5 | 4.1 |
+
+
+ | w/o SRC-MoE |
+ 19.6 | - | 32.5 | 30.9 | 24.4 | 24.5 | 14.3 | 28.7 | 25.8 |
+
+
+
+
+
+ | es |
+ HENT-SRT-M20×9 |
+ 9.9 | 22.1 | - | 20.2 | 15.7 | 15.1 | 6.9 | 22.4 | 12.2 |
+
+
+ | HENT-SRT-M2M |
+ 2.1 | 13.4 | - | 3.9 | 1.5 | 3.1 | 0.9 | 5.4 | 2.2 |
+
+
+ | LCMA-SRT |
+ 13.7 | 26.1 | - | 26.3 | 21.0 | 19.4 | 10.3 | 26.6 | 17.7 |
+
+
+ | TGT-MoE→MoE |
+ 2.0 | 13.9 | - | 3.5 | 1.6 | 2.3 | 1.1 | 5.0 | 1.9 |
+
+
+ | TGT-MoE→T-Bias |
+ 11.8 | 23.3 | - | 23.7 | 18.4 | 17.3 | 7.7 | 23.7 | 14.8 |
+
+
+ | w/o TGT-MoE |
+ 1.8 | 13.2 | - | 3.7 | 1.4 | 2.5 | 1.1 | 5.5 | 2.2 |
+
+
+ | w/o SRC-MoE |
+ 13.3 | 25.1 | - | 26.5 | 20.6 | 19.3 | 10.2 | 26.1 | 17.5 |
+
+
+
+
+
+ | fr |
+ HENT-SRT-M20×9 |
+ 11.0 | 23.5 | 20.3 | - | 17.6 | 16.9 | 7.4 | 23.3 | 13.0 |
+
+
+ | HENT-SRT-M2M |
+ 2.9 | 11.9 | 6.4 | - | 2.2 | 4.0 | 1.3 | 6.5 | 2.4 |
+
+
+ | LCMA-SRT |
+ 14.9 | 28.6 | 27.0 | - | 22.5 | 21.3 | 11.1 | 27.5 | 18.3 |
+
+
+ | TGT-MoE→MoE |
+ 2.6 | 13.8 | 5.2 | - | 2.0 | 3.0 | 1.3 | 5.9 | 2.1 |
+
+
+ | TGT-MoE→T-Bias |
+ 13.3 | 24.8 | 24.7 | - | 20.1 | 18.5 | 9.0 | 25.0 | 15.2 |
+
+
+ | w/o TGT-MoE |
+ 2.3 | 11.8 | 6.8 | - | 2.2 | 3.4 | 1.5 | 6.1 | 2.3 |
+
+
+ | w/o SRC-MoE |
+ 14.1 | 27.3 | 25.9 | - | 22.4 | 20.2 | 10.9 | 26.9 | 18.6 |
+
+
+
+
+
+ | it |
+ HENT-SRT-M20×9 |
+ 11.3 | 23.0 | 21.3 | 20.3 | - | 16.1 | 8.3 | 22.4 | 13.4 |
+
+
+ | HENT-SRT-M2M |
+ 2.9 | 14.7 | 5.1 | 4.0 | - | 3.2 | 1.7 | 5.6 | 2.0 |
+
+
+ | LCMA-SRT |
+ 14.8 | 27.0 | 27.3 | 25.3 | - | 20.2 | 11.0 | 26.1 | 17.8 |
+
+
+ | TGT-MoE→MoE |
+ 2.6 | 16.8 | 4.3 | 3.0 | - | 2.6 | 1.6 | 5.2 | 1.7 |
+
+
+ | TGT-MoE→T-Bias |
+ 12.9 | 23.5 | 24.9 | 22.7 | - | 17.3 | 8.8 | 24.1 | 14.2 |
+
+
+ | w/o TGT-MoE |
+ 2.1 | 15.2 | 5.1 | 3.5 | - | 2.8 | 1.7 | 5.4 | 1.8 |
+
+
+ | w/o SRC-MoE |
+ 14.0 | 26.0 | 26.8 | 25.1 | - | 19.4 | 11.2 | 26.5 | 17.7 |
+
+
+
+
+
+ | nl |
+ HENT-SRT-M20×9 |
+ 7.1 | 15.6 | 11.3 | 10.4 | 7.3 | - | 3.7 | 10.4 | 6.3 |
+
+
+ | HENT-SRT-M2M |
+ 2.3 | 9.8 | 3.1 | 2.6 | 1.2 | - | 0.9 | 2.9 | 1.9 |
+
+
+ | LCMA-SRT |
+ 12.1 | 21.0 | 17.6 | 16.5 | 13.6 | - | 7.0 | 16.9 | 11.6 |
+
+
+ | TGT-MoE→MoE |
+ 1.7 | 12.0 | 2.1 | 1.8 | 0.9 | - | 0.6 | 2.1 | 1.2 |
+
+
+ | TGT-MoE→T-Bias |
+ 10.1 | 18.1 | 16.3 | 15.3 | 11.7 | - | 4.9 | 15.5 | 10.0 |
+
+
+ | w/o TGT-MoE |
+ 1.6 | 10.8 | 2.8 | 2.6 | 1.0 | - | 1.1 | 2.7 | 1.4 |
+
+
+ | w/o SRC-MoE |
+ 11.8 | 20.0 | 17.5 | 16.8 | 12.8 | - | 6.7 | 17.0 | 11.8 |
+
+
+
+
+
+ | pl |
+ HENT-SRT-M20×9 |
+ 9.5 | 19.3 | 17.1 | 15.7 | 11.9 | 14.3 | - | 14.6 | 10.0 |
+
+
+ | HENT-SRT-M2M |
+ 2.4 | 12.1 | 4.6 | 3.8 | 1.6 | 3.4 | - | 3.8 | 2.1 |
+
+
+ | LCMA-SRT |
+ 14.3 | 23.9 | 24.1 | 22.9 | 18.6 | 19.5 | - | 20.8 | 16.5 |
+
+
+ | TGT-MoE→MoE |
+ 2.2 | 13.9 | 3.7 | 3.0 | 1.5 | 2.2 | - | 3.2 | 1.4 |
+
+
+ | TGT-MoE→T-Bias |
+ 12.3 | 21.5 | 22.1 | 21.1 | 16.5 | 17.6 | - | 19.4 | 13.4 |
+
+
+ | w/o TGT-MoE |
+ 2.1 | 11.6 | 4.9 | 4.2 | 1.3 | 2.9 | - | 4.0 | 2.0 |
+
+
+ | w/o SRC-MoE |
+ 13.5 | 23.3 | 22.6 | 22.3 | 18.0 | 19.2 | - | 20.9 | 16.2 |
+
+
+
+
+
+ | pt |
+ HENT-SRT-M20×9 |
+ 10.9 | 23.7 | 22.1 | 21.3 | 17.3 | 15.6 | 7.5 | - | 13.9 |
+
+
+ | HENT-SRT-M2M |
+ 2.3 | 13.3 | 6.6 | 4.0 | 1.9 | 3.0 | 1.1 | - | 2.5 |
+
+
+ | LCMA-SRT |
+ 15.4 | 28.1 | 28.3 | 27.0 | 22.8 | 19.7 | 10.5 | - | 19.0 |
+
+
+ | TGT-MoE→MoE |
+ 1.9 | 17.4 | 4.7 | 3.4 | 1.4 | 2.1 | 1.0 | - | 1.7 |
+
+
+ | TGT-MoE→T-Bias |
+ 13.6 | 24.9 | 25.6 | 24.9 | 20.7 | 18.0 | 8.8 | - | 16.1 |
+
+
+ | w/o TGT-MoE |
+ 1.7 | 14.6 | 6.7 | 3.9 | 1.4 | 2.3 | 1.2 | - | 1.8 |
+
+
+ | w/o SRC-MoE |
+ 14.5 | 26.7 | 27.5 | 27.1 | 22.5 | 19.0 | 10.6 | - | 19.0 |
+
+
+
+
+
+ | ro |
+ HENT-SRT-M20×9 |
+ 10.9 | 25.3 | 21.4 | 21.4 | 15.8 | 16.0 | 7.9 | 18.8 | - |
+
+
+ | HENT-SRT-M2M |
+ 1.9 | 16.4 | 4.6 | 3.7 | 1.5 | 2.2 | 0.7 | 3.9 | - |
+
+
+ | LCMA-SRT |
+ 15.8 | 30.1 | 28.9 | 28.4 | 22.1 | 19.7 | 12.2 | 25.3 | - |
+
+
+ | TGT-MoE→MoE |
+ 1.9 | 17.5 | 4.1 | 3.4 | 1.6 | 1.8 | 0.7 | 3.4 | - |
+
+
+ | TGT-MoE→T-Bias |
+ 13.7 | 25.9 | 26.3 | 25.4 | 19.6 | 17.6 | 9.1 | 23.3 | - |
+
+
+ | w/o TGT-MoE |
+ 1.5 | 13.5 | 6.5 | 3.9 | 1.5 | 2.3 | 0.9 | 4.5 | - |
+
+
+ | w/o SRC-MoE |
+ 15.0 | 29.2 | 27.7 | 28.3 | 21.8 | 19.2 | 11.8 | 24.7 | - |
+
+
+
+
+### COMET
+
+
+
+
+ | SRC\TGT |
+ Model |
+ COMET ↑ |
+
+
+ | de | en | es | fr | it | nl | pl | pt | ro |
+
+
+
+
+
+ | de |
+ HENT-SRT-M20×9 |
+ - | 0.615 | 0.531 | 0.479 | 0.504 | 0.549 | 0.521 | 0.544 | 0.545 |
+
+
+ | HENT-SRT-M2M |
+ - | 0.522 | 0.453 | 0.407 | 0.409 | 0.397 | 0.383 | 0.447 | 0.391 |
+
+
+ | LCMA-SRT |
+ - | 0.683 | 0.624 | 0.572 | 0.591 | 0.604 | 0.591 | 0.636 | 0.627 |
+
+
+ | TGT-MoE→MoE |
+ - | 0.528 | 0.451 | 0.400 | 0.412 | 0.383 | 0.377 | 0.446 | 0.393 |
+
+
+ | TGT-MoE→T-Bias |
+ - | 0.640 | 0.590 | 0.533 | 0.557 | 0.546 | 0.542 | 0.601 | 0.572 |
+
+
+ | w/o TGT-MoE |
+ - | 0.524 | 0.454 | 0.407 | 0.408 | 0.393 | 0.381 | 0.448 | 0.391 |
+
+
+ | w/o SRC-MoE |
+ - | 0.675 | 0.614 | 0.565 | 0.578 | 0.601 | 0.578 | 0.629 | 0.626 |
+
+
+
+
+
+ | en |
+ HENT-SRT-M20×9 |
+ 0.571 | - | 0.641 | 0.606 | 0.625 | 0.620 | 0.584 | 0.668 | 0.680 |
+
+
+ | HENT-SRT-M2M |
+ 0.421 | - | 0.533 | 0.470 | 0.487 | 0.430 | 0.419 | 0.524 | 0.458 |
+
+
+ | LCMA-SRT |
+ 0.638 | - | 0.741 | 0.690 | 0.714 | 0.674 | 0.663 | 0.749 | 0.765 |
+
+
+ | TGT-MoE→MoE |
+ 0.422 | - | 0.535 | 0.467 | 0.495 | 0.424 | 0.421 | 0.528 | 0.463 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.582 | - | 0.696 | 0.637 | 0.664 | 0.612 | 0.590 | 0.704 | 0.693 |
+
+
+ | w/o TGT-MoE |
+ 0.413 | - | 0.540 | 0.469 | 0.488 | 0.428 | 0.423 | 0.532 | 0.459 |
+
+
+ | w/o SRC-MoE |
+ 0.626 | - | 0.733 | 0.690 | 0.705 | 0.661 | 0.649 | 0.742 | 0.762 |
+
+
+
+
+
+ | es |
+ HENT-SRT-M20×9 |
+ 0.488 | 0.652 | - | 0.548 | 0.571 | 0.534 | 0.546 | 0.636 | 0.589 |
+
+
+ | HENT-SRT-M2M |
+ 0.357 | 0.536 | - | 0.416 | 0.424 | 0.385 | 0.374 | 0.464 | 0.396 |
+
+
+ | LCMA-SRT |
+ 0.544 | 0.708 | - | 0.627 | 0.657 | 0.584 | 0.609 | 0.709 | 0.663 |
+
+
+ | TGT-MoE→MoE |
+ 0.361 | 0.541 | - | 0.416 | 0.430 | 0.377 | 0.375 | 0.468 | 0.394 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.503 | 0.669 | - | 0.579 | 0.609 | 0.539 | 0.556 | 0.665 | 0.608 |
+
+
+ | w/o TGT-MoE |
+ 0.358 | 0.539 | - | 0.419 | 0.425 | 0.382 | 0.376 | 0.466 | 0.394 |
+
+
+ | w/o SRC-MoE |
+ 0.541 | 0.702 | - | 0.622 | 0.651 | 0.578 | 0.600 | 0.705 | 0.666 |
+
+
+
+
+
+ | fr |
+ HENT-SRT-M20×9 |
+ 0.499 | 0.685 | 0.603 | - | 0.603 | 0.551 | 0.555 | 0.650 | 0.618 |
+
+
+ | HENT-SRT-M2M |
+ 0.373 | 0.535 | 0.484 | - | 0.440 | 0.396 | 0.385 | 0.481 | 0.408 |
+
+
+ | LCMA-SRT |
+ 0.561 | 0.737 | 0.700 | - | 0.685 | 0.603 | 0.616 | 0.723 | 0.701 |
+
+
+ | TGT-MoE→MoE |
+ 0.379 | 0.552 | 0.482 | - | 0.442 | 0.391 | 0.385 | 0.482 | 0.410 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.519 | 0.696 | 0.658 | - | 0.636 | 0.555 | 0.566 | 0.674 | 0.637 |
+
+
+ | w/o TGT-MoE |
+ 0.372 | 0.538 | 0.485 | - | 0.440 | 0.393 | 0.385 | 0.484 | 0.411 |
+
+
+ | w/o SRC-MoE |
+ 0.553 | 0.730 | 0.685 | - | 0.673 | 0.592 | 0.605 | 0.711 | 0.695 |
+
+
+
+
+
+ | it |
+ HENT-SRT-M20×9 |
+ 0.507 | 0.679 | 0.614 | 0.569 | - | 0.551 | 0.568 | 0.650 | 0.623 |
+
+
+ | HENT-SRT-M2M |
+ 0.372 | 0.560 | 0.477 | 0.425 | - | 0.393 | 0.380 | 0.472 | 0.404 |
+
+
+ | LCMA-SRT |
+ 0.560 | 0.728 | 0.698 | 0.640 | - | 0.600 | 0.619 | 0.717 | 0.686 |
+
+
+ | TGT-MoE→MoE |
+ 0.374 | 0.578 | 0.476 | 0.428 | - | 0.391 | 0.385 | 0.478 | 0.412 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.520 | 0.689 | 0.657 | 0.593 | - | 0.553 | 0.570 | 0.672 | 0.630 |
+
+
+ | w/o TGT-MoE |
+ 0.370 | 0.572 | 0.483 | 0.428 | - | 0.393 | 0.384 | 0.476 | 0.411 |
+
+
+ | w/o SRC-MoE |
+ 0.558 | 0.722 | 0.693 | 0.635 | - | 0.591 | 0.615 | 0.711 | 0.689 |
+
+
+
+
+
+ | nl |
+ HENT-SRT-M20×9 |
+ 0.444 | 0.581 | 0.500 | 0.460 | 0.467 | - | 0.486 | 0.509 | 0.509 |
+
+
+ | HENT-SRT-M2M |
+ 0.367 | 0.508 | 0.435 | 0.397 | 0.402 | - | 0.365 | 0.435 | 0.380 |
+
+
+ | LCMA-SRT |
+ 0.538 | 0.660 | 0.595 | 0.544 | 0.561 | - | 0.556 | 0.604 | 0.593 |
+
+
+ | TGT-MoE→MoE |
+ 0.359 | 0.532 | 0.431 | 0.395 | 0.403 | - | 0.367 | 0.432 | 0.378 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.493 | 0.615 | 0.556 | 0.510 | 0.527 | - | 0.512 | 0.569 | 0.549 |
+
+
+ | w/o TGT-MoE |
+ 0.361 | 0.517 | 0.434 | 0.398 | 0.400 | - | 0.372 | 0.433 | 0.377 |
+
+
+ | w/o SRC-MoE |
+ 0.528 | 0.651 | 0.582 | 0.538 | 0.550 | - | 0.545 | 0.601 | 0.587 |
+
+
+
+
+
+ | pl |
+ HENT-SRT-M20×9 |
+ 0.515 | 0.643 | 0.568 | 0.518 | 0.545 | 0.543 | - | 0.584 | 0.583 |
+
+
+ | HENT-SRT-M2M |
+ 0.385 | 0.539 | 0.469 | 0.424 | 0.429 | 0.397 | - | 0.462 | 0.401 |
+
+
+ | LCMA-SRT |
+ 0.584 | 0.709 | 0.667 | 0.612 | 0.651 | 0.608 | - | 0.683 | 0.677 |
+
+
+ | TGT-MoE→MoE |
+ 0.380 | 0.552 | 0.464 | 0.419 | 0.431 | 0.388 | - | 0.462 | 0.402 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.539 | 0.669 | 0.633 | 0.572 | 0.605 | 0.561 | - | 0.645 | 0.624 |
+
+
+ | w/o TGT-MoE |
+ 0.379 | 0.531 | 0.466 | 0.421 | 0.427 | 0.395 | - | 0.461 | 0.402 |
+
+
+ | w/o SRC-MoE |
+ 0.575 | 0.698 | 0.655 | 0.596 | 0.634 | 0.601 | - | 0.670 | 0.669 |
+
+
+
+
+
+ | pt |
+ HENT-SRT-M20×9 |
+ 0.522 | 0.692 | 0.631 | 0.584 | 0.605 | 0.556 | 0.576 | - | 0.636 |
+
+
+ | HENT-SRT-M2M |
+ 0.381 | 0.557 | 0.491 | 0.439 | 0.444 | 0.402 | 0.390 | - | 0.409 |
+
+
+ | LCMA-SRT |
+ 0.581 | 0.744 | 0.722 | 0.662 | 0.695 | 0.609 | 0.632 | - | 0.710 |
+
+
+ | TGT-MoE→MoE |
+ 0.381 | 0.594 | 0.482 | 0.436 | 0.444 | 0.397 | 0.392 | - | 0.409 |
+
+
+ | TGT-MoE→T-Bias |
+ 0.541 | 0.708 | 0.684 | 0.623 | 0.654 | 0.566 | 0.585 | - | 0.654 |
+
+
+ | w/o TGT-MoE |
+ 0.377 | 0.572 | 0.489 | 0.438 | 0.440 | 0.396 | 0.387 | - | 0.409 |
+
+
+ | w/o SRC-MoE |
+ 0.577 | 0.738 | 0.710 | 0.656 | 0.688 | 0.604 | 0.624 | - | 0.708 |
+
+
+
+
+
+ | ro |
+ HENT-SRT-M20×9 |
+ 0.514 | 0.697 | 0.606 | 0.575 | 0.596 | 0.563 | 0.569 | 0.627 | - |
+
+
+ | HENT-SRT-M2M |
+ 0.381 | 0.585 | 0.487 | 0.443 | 0.446 | 0.408 | 0.386 | 0.480 | - |
+
+
+ | LCMA-SRT |
+ 0.587 | 0.753 | 0.711 | 0.667 | 0.696 | 0.624 | 0.642 | 0.724 | - |
+
+
+ | TGT-MoE→MoE |
+ 0.384 | 0.591 | 0.486 | 0.443 | 0.449 | 0.405 | 0.388 | 0.484 | - |
+
+
+ | TGT-MoE→T-Bias |
+ 0.537 | 0.713 | 0.665 | 0.616 | 0.644 | 0.574 | 0.574 | 0.680 | - |
+
+
+ | w/o TGT-MoE |
+ 0.377 | 0.563 | 0.488 | 0.441 | 0.447 | 0.405 | 0.386 | 0.480 | - |
+
+
+ | w/o SRC-MoE |
+ 0.582 | 0.747 | 0.699 | 0.662 | 0.687 | 0.618 | 0.628 | 0.714 | - |
+
+
+
+
+### LMR
+
+ | SRC\TGT | Model | LMR (%) ↓ |
| de | en | es | fr | it | nl | pl | pt | ro |
| de | HENT-SRT-M20×9 | - | 0.08 | 0.70 | 0.64 | 0.66 | 1.00 | 0.00 | 3.39 | 1.70 |
| HENT-SRT-M2M | - | 56.60 | 87.83 | 90.14 | 94.98 | 77.09 | 83.11 | 83.75 | 82.55 |
| LCMA-SRT | - | 0.38 | 0.49 | 0.43 | 0.58 | 0.84 | 0.73 | 2.38 | 1.79 |
| TGT-MoE→MoE | - | 44.09 | 91.70 | 93.43 | 96.22 | 87.51 | 88.57 | 86.63 | 84.42 |
| TGT-MoE→T-Bias | - | 0.34 | 0.56 | 0.64 | 0.82 | 1.84 | 0.36 | 2.89 | 1.62 |
| w/o TGT-MoE | - | 54.96 | 87.83 | 91.79 | 97.04 | 83.22 | 80.86 | 83.02 | 83.12 |
| w/o SRC-MoE | - | 0.38 | 0.84 | 0.86 | 1.48 | 0.92 | 0.58 | 3.10 | 1.62 |
| en | HENT-SRT-M20×9 | 0.00 | - | 0.79 | 0.25 | 0.35 | 0.65 | 0.16 | 1.98 | 1.64 |
| HENT-SRT-M2M | 78.21 | - | 78.93 | 86.08 | 95.93 | 78.54 | 85.14 | 81.62 | 80.55 |
| LCMA-SRT | 0.08 | - | 0.95 | 0.16 | 0.09 | 0.24 | 0.40 | 1.98 | 1.64 |
| TGT-MoE→MoE | 79.25 | - | 82.40 | 85.26 | 95.75 | 81.54 | 87.16 | 79.24 | 79.45 |
| TGT-MoE→T-Bias | 0.32 | - | 0.71 | 0.25 | 0.27 | 0.49 | 0.32 | 2.38 | 1.28 |
| w/o TGT-MoE | 83.56 | - | 77.82 | 86.24 | 95.39 | 83.48 | 84.49 | 78.76 | 81.74 |
| w/o SRC-MoE | 0.08 | - | 0.55 | 0.08 | 0.53 | 0.40 | 0.32 | 2.38 | 1.83 |
| es | HENT-SRT-M20×9 | 0.18 | 0.22 | - | 0.18 | 0.46 | 1.01 | 0.09 | 1.19 | 0.88 |
| HENT-SRT-M2M | 88.51 | 58.54 | - | 92.98 | 96.76 | 89.58 | 90.93 | 80.07 | 86.81 |
| LCMA-SRT | 0.54 | 0.61 | - | 0.37 | 0.65 | 0.92 | 0.38 | 1.29 | 1.54 |
| TGT-MoE→MoE | 88.87 | 50.66 | - | 93.16 | 96.85 | 92.59 | 92.63 | 82.09 | 88.68 |
| TGT-MoE→T-Bias | 0.36 | 0.61 | - | 0.28 | 0.37 | 1.28 | 0.19 | 1.47 | 1.10 |
| w/o TGT-MoE | 91.56 | 55.95 | - | 92.79 | 97.59 | 91.58 | 87.43 | 76.19 | 85.71 |
| w/o SRC-MoE | 0.27 | 0.50 | - | 0.55 | 0.74 | 0.82 | 0.47 | 1.47 | 1.43 |
| fr | HENT-SRT-M20×9 | 0.00 | 0.11 | 0.55 | - | 0.00 | 0.26 | 0.18 | 1.55 | 1.16 |
| HENT-SRT-M2M | 85.18 | 71.51 | 87.89 | - | 95.60 | 86.66 | 89.85 | 77.71 | 84.93 |
| LCMA-SRT | 0.00 | 0.33 | 0.73 | - | 0.38 | 0.87 | 0.36 | 1.91 | 1.16 |
| TGT-MoE→MoE | 88.27 | 57.97 | 90.88 | - | 96.84 | 90.66 | 90.82 | 81.88 | 85.76 |
| TGT-MoE→T-Bias | 0.18 | 0.39 | 0.36 | - | 0.19 | 1.31 | 0.27 | 0.91 | 1.37 |
| w/o TGT-MoE | 88.91 | 70.85 | 87.41 | - | 96.46 | 89.61 | 88.84 | 78.05 | 83.02 |
| w/o SRC-MoE | 0.37 | 0.44 | 0.46 | - | 0.57 | 0.61 | 0.27 | 1.18 | 1.90 |
| it | HENT-SRT-M20×9 | 0.11 | 0.00 | 0.57 | 0.23 | - | 0.36 | 0.25 | 2.20 | 0.54 |
| HENT-SRT-M2M | 88.54 | 57.23 | 92.05 | 94.25 | - | 91.61 | 92.77 | 86.59 | 90.38 |
| LCMA-SRT | 0.11 | 0.42 | 0.91 | 0.23 | - | 0.84 | 0.37 | 1.62 | 0.95 |
| TGT-MoE→MoE | 91.92 | 41.43 | 91.37 | 94.13 | - | 94.36 | 93.87 | 87.40 | 89.70 |
| TGT-MoE→T-Bias | 0.44 | 0.12 | 0.45 | 0.45 | - | 0.72 | 0.00 | 1.73 | 1.49 |
| w/o TGT-MoE | 92.36 | 47.80 | 91.37 | 95.94 | - | 93.04 | 91.18 | 87.51 | 89.30 |
| w/o SRC-MoE | 0.00 | 0.12 | 0.45 | 0.79 | - | 0.72 | 0.12 | 1.50 | 0.81 |
| nl | HENT-SRT-M20×9 | 0.09 | 0.34 | 0.49 | 0.49 | 0.90 | - | 0.21 | 3.18 | 1.60 |
| HENT-SRT-M2M | 78.25 | 56.62 | 88.85 | 89.02 | 93.71 | - | 86.96 | 82.91 | 81.98 |
| LCMA-SRT | 0.19 | 0.86 | 1.19 | 0.79 | 0.79 | - | 0.62 | 1.80 | 1.94 |
| TGT-MoE→MoE | 82.38 | 35.99 | 94.86 | 92.08 | 95.38 | - | 89.00 | 86.94 | 86.99 |
| TGT-MoE→T-Bias | 0.47 | 0.52 | 1.38 | 0.69 | 0.79 | - | 0.83 | 4.03 | 1.71 |
| w/o TGT-MoE | 83.47 | 49.05 | 88.42 | 89.30 | 95.72 | - | 81.93 | 84.82 | 82.65 |
| w/o SRC-MoE | 0.19 | 0.63 | 0.99 | 1.58 | 2.02 | - | 1.04 | 2.23 | 2.17 |
| pl | HENT-SRT-M20×9 | 0.00 | 0.22 | 0.72 | 0.16 | 0.42 | 1.80 | - | 2.08 | 1.72 |
| HENT-SRT-M2M | 82.54 | 60.55 | 89.70 | 91.34 | 96.01 | 86.68 | - | 83.45 | 87.29 |
| LCMA-SRT | 0.00 | 0.40 | 0.80 | 0.48 | 0.85 | 1.14 | - | 1.60 | 1.01 |
| TGT-MoE→MoE | 85.89 | 46.27 | 92.66 | 92.05 | 96.94 | 90.28 | - | 88.57 | 88.70 |
| TGT-MoE→T-Bias | 0.39 | 0.49 | 1.04 | 0.32 | 0.51 | 1.47 | - | 1.44 | 1.31 |
| w/o TGT-MoE | 86.52 | 64.23 | 89.31 | 91.49 | 97.28 | 89.13 | - | 82.57 | 86.28 |
| w/o SRC-MoE | 0.00 | 0.36 | 0.40 | 0.79 | 0.85 | 1.39 | - | 1.60 | 1.21 |
| pt | HENT-SRT-M20×9 | 0.00 | 0.17 | 0.16 | 0.16 | 0.08 | 0.41 | 0.08 | - | 0.81 |
| HENT-SRT-M2M | 88.36 | 65.79 | 86.39 | 91.99 | 97.51 | 90.31 | 90.80 | - | 87.09 |
| LCMA-SRT | 0.16 | 0.39 | 0.08 | 0.24 | 0.08 | 0.81 | 0.50 | - | 1.17 |
| TGT-MoE→MoE | 90.32 | 40.07 | 92.60 | 94.50 | 97.68 | 94.79 | 92.56 | - | 91.43 |
| TGT-MoE→T-Bias | 0.24 | 0.48 | 0.24 | 0.24 | 0.00 | 0.90 | 0.08 | - | 0.81 |
| w/o TGT-MoE | 91.50 | 54.86 | 85.19 | 92.93 | 98.76 | 93.65 | 90.22 | - | 89.89 |
| w/o SRC-MoE | 0.00 | 0.26 | 0.08 | 0.47 | 0.17 | 0.49 | 0.25 | - | 1.08 |
| ro | HENT-SRT-M20×9 | 0.00 | 0.15 | 0.42 | 0.26 | 0.77 | 0.91 | 0.00 | 1.58 | - |
| HENT-SRT-M2M | 90.98 | 50.74 | 93.27 | 95.33 | 98.46 | 93.14 | 93.04 | 89.42 | - |
| LCMA-SRT | 0.16 | 0.56 | 0.50 | 0.17 | 0.77 | 1.24 | 0.43 | 1.33 | - |
| TGT-MoE→MoE | 91.39 | 45.21 | 93.69 | 94.99 | 97.52 | 94.38 | 94.93 | 89.92 | - |
| TGT-MoE→T-Bias | 0.08 | 0.31 | 0.42 | 0.17 | 0.34 | 0.83 | 0.17 | 2.00 | - |
| w/o TGT-MoE | 92.61 | 65.55 | 87.96 | 94.64 | 98.80 | 92.81 | 91.92 | 86.75 | - |
| w/o SRC-MoE | 0.16 | 0.31 | 0.58 | 0.26 | 1.03 | 0.41 | 0.60 | 2.00 | - |
diff --git a/egs/europarl_st/SRT/lcma_srt/LCMA-SRT.png b/egs/europarl_st/SRT/lcma_srt/LCMA-SRT.png
new file mode 100644
index 0000000000..0c2fb3a93c
Binary files /dev/null and b/egs/europarl_st/SRT/lcma_srt/LCMA-SRT.png differ
diff --git a/egs/europarl_st/SRT/lcma_srt/README.md b/egs/europarl_st/SRT/lcma_srt/README.md
new file mode 100644
index 0000000000..6b3812247c
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/README.md
@@ -0,0 +1,74 @@
+# LCMA-SRT: Language-Conditional Mixture-of-Experts Adapters for Joint Multilingual Speech Recognition and Translation
+
+This recipe implements LCMA-SRT on the [Europarl-ST](https://www.mllp.upv.es/europarl-st/) dataset (9 languages, 72 translation directions).
+
+Paper: [ACL 2026](https://aclanthology.org/2026.acl-long.1634/)
+
+## Model Architecture
+
+LCMA-SRT augments a hierarchical transducer (Zipformer encoder) with:
+- **SRC-MoE adapter**: Source-conditioned Mixture-of-Experts adapter after the ASR encoder to reduce cross-language interference
+- **TGT-MoE adapter**: Target-conditioned Mixture-of-Experts adapter for the ST encoder to stabilize target-language generation
+
+## Performance
+
+| Model | WER (%) ↓ | Avg BLEU ↑ | Avg COMET ↑ | LMR (%) ↓ |
+|-------|-----------|-----------|-------------|-----------|
+| HENT-SRT-M2O×9 | 23.28 | 15.3 | 0.575 | 0.65 |
+| HENT-SRT-M2M | 16.65 | 4.3 | 0.436 | 84.95 |
+| **LCMA-SRT** | **15.71** | **20.5** | **0.651** | **0.75** |
+
+## Usage
+
+### Data Preparation
+
+```bash
+cd egs/europarl_st/SRT
+bash prepare.sh --stage 0 --stop-stage 5
+```
+
+See [local/README.md](../local/README.md) for detailed documentation of each preprocessing step.
+
+### Training
+
+#### Stage 1: Multilingual ASR Pretraining (with SRC-MoE)
+
+```bash
+bash lcma_srt/train/stage1/cr_ctc_sc_moe.sh
+```
+
+#### Stage 2: Many-to-Many Joint ASR+ST Training
+
+```bash
+bash lcma_srt/train/stage2/lcma_srt.sh
+```
+
+### Decoding
+
+#### Stage 1: ASR Decoding
+
+```bash
+bash lcma_srt/decode/stage1/decode_cr_ctc_sc_moe.sh
+```
+
+#### Stage 2: Joint ASR+ST Decoding
+
+```bash
+bash lcma_srt/decode/stage2/decode_lcma_srt.sh
+```
+
+## Pre-trained Models
+
+Pre-trained checkpoints are available on [Hugging Face](https://huggingface.co/linanjie0820/lcma-srt).
+
+## Citation
+
+```bibtex
+@inproceedings{li2026lcma,
+ title={LCMA-SRT: Language-Conditional Mixture-of-Experts Adapters for Joint Multilingual Speech Recognition and Translation},
+ author={Li, Nanjie and Guo, Xiaoyong and Huang, Hao and Haihua, Xu and Shi, Wei},
+ booktitle={Proceedings of the 64th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)},
+ pages={35363--35377},
+ year={2026}
+}
+```
diff --git a/egs/europarl_st/SRT/lcma_srt/attention_decoder.py b/egs/europarl_st/SRT/lcma_srt/attention_decoder.py
new file mode 100644
index 0000000000..bff536f90b
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/attention_decoder.py
@@ -0,0 +1,583 @@
+#!/usr/bin/env python3
+# Copyright 2023 Xiaomi Corp. (authors: Zengwei Yao)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import math
+from typing import List, Optional
+
+import k2
+import torch
+import torch.nn as nn
+from label_smoothing import LabelSmoothingLoss
+from scaling import penalize_abs_values_gt
+
+from icefall.utils import add_eos, add_sos, make_pad_mask
+
+
+class AttentionDecoderModel(nn.Module):
+ """
+ Args:
+ vocab_size (int): Number of classes.
+ decoder_dim: (int,int): embedding dimension of 2 encoder stacks
+ attention_dim: (int,int): attention dimension of 2 encoder stacks
+ num_heads (int, int): number of heads
+ dim_feedforward (int, int): feedforward dimension in 2 encoder stacks
+ num_encoder_layers (int): number of encoder layers
+ dropout (float): dropout rate
+ """
+
+ def __init__(
+ self,
+ vocab_size: int,
+ decoder_dim: int = 512,
+ num_decoder_layers: int = 6,
+ attention_dim: int = 512,
+ num_heads: int = 8,
+ feedforward_dim: int = 2048,
+ memory_dim: int = 512,
+ sos_id: int = 1,
+ eos_id: int = 1,
+ dropout: float = 0.1,
+ ignore_id: int = -1,
+ label_smoothing: float = 0.1,
+ ):
+ super().__init__()
+ self.eos_id = eos_id
+ self.sos_id = sos_id
+ self.ignore_id = ignore_id
+
+ # For the segment of the warmup period, we let the Embedding
+ # layer learn something. Then we start to warm up the other encoders.
+ self.decoder = TransformerDecoder(
+ vocab_size=vocab_size,
+ d_model=decoder_dim,
+ num_decoder_layers=num_decoder_layers,
+ attention_dim=attention_dim,
+ num_heads=num_heads,
+ feedforward_dim=feedforward_dim,
+ memory_dim=memory_dim,
+ dropout=dropout,
+ )
+
+ # Used to calculate attention-decoder loss
+ self.loss_fun = LabelSmoothingLoss(
+ ignore_index=ignore_id, label_smoothing=label_smoothing, reduction="sum"
+ )
+
+ def _pre_ys_in_out(self, ys: k2.RaggedTensor, ys_lens: torch.Tensor):
+ """Prepare ys_in_pad and ys_out_pad."""
+ ys_in = add_sos(ys, sos_id=self.sos_id)
+ # [B, S+1], start with SOS
+ ys_in_pad = ys_in.pad(mode="constant", padding_value=self.eos_id)
+ ys_in_lens = ys_lens + 1
+
+ ys_out = add_eos(ys, eos_id=self.eos_id)
+ # [B, S+1], end with EOS
+ ys_out_pad = ys_out.pad(mode="constant", padding_value=self.ignore_id)
+
+ return ys_in_pad.to(torch.int64), ys_in_lens, ys_out_pad.to(torch.int64)
+
+ def calc_att_loss(
+ self,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ ys: k2.RaggedTensor,
+ ys_lens: torch.Tensor,
+ ) -> torch.Tensor:
+ """Calculate attention-decoder loss.
+ Args:
+ encoder_out: (batch, num_frames, encoder_dim)
+ encoder_out_lens: (batch,)
+ token_ids: A list of token id list.
+
+ Return: The attention-decoder loss.
+ """
+ ys_in_pad, ys_in_lens, ys_out_pad = self._pre_ys_in_out(ys, ys_lens)
+
+ # decoder forward
+ decoder_out = self.decoder(
+ x=ys_in_pad,
+ x_lens=ys_in_lens,
+ memory=encoder_out,
+ memory_lens=encoder_out_lens,
+ )
+
+ loss = self.loss_fun(x=decoder_out, target=ys_out_pad)
+ return loss
+
+ def nll(
+ self,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ token_ids: List[List[int]],
+ ) -> torch.Tensor:
+ """Compute negative log likelihood(nll) from attention-decoder.
+ Args:
+ encoder_out: (batch, num_frames, encoder_dim)
+ encoder_out_lens: (batch,)
+ token_ids: A list of token id list.
+
+ Return: A tensor of shape (batch, num_tokens).
+ """
+ ys = k2.RaggedTensor(token_ids).to(device=encoder_out.device)
+ row_splits = ys.shape.row_splits(1)
+ ys_lens = row_splits[1:] - row_splits[:-1]
+
+ ys_in_pad, ys_in_lens, ys_out_pad = self._pre_ys_in_out(ys, ys_lens)
+
+ # decoder forward
+ decoder_out = self.decoder(
+ x=ys_in_pad,
+ x_lens=ys_in_lens,
+ memory=encoder_out,
+ memory_lens=encoder_out_lens,
+ )
+
+ batch_size, _, num_classes = decoder_out.size()
+ nll = nn.functional.cross_entropy(
+ decoder_out.view(-1, num_classes),
+ ys_out_pad.view(-1),
+ ignore_index=self.ignore_id,
+ reduction="none",
+ )
+ nll = nll.view(batch_size, -1)
+ return nll
+
+
+class TransformerDecoder(nn.Module):
+ """Transfomer decoder module.
+
+ Args:
+ vocab_size: output dim
+ d_model: decoder dimension
+ num_decoder_layers: number of decoder layers
+ attention_dim: total dimension of multi head attention
+ num_heads: number of attention heads
+ feedforward_dim: hidden dimension of feed_forward module
+ dropout: dropout rate
+ """
+
+ def __init__(
+ self,
+ vocab_size: int,
+ d_model: int = 512,
+ num_decoder_layers: int = 6,
+ attention_dim: int = 512,
+ num_heads: int = 8,
+ feedforward_dim: int = 2048,
+ memory_dim: int = 512,
+ dropout: float = 0.1,
+ ):
+ super().__init__()
+ self.embed = nn.Embedding(num_embeddings=vocab_size, embedding_dim=d_model)
+
+ # Absolute positional encoding
+ self.pos = PositionalEncoding(d_model, dropout_rate=0.1)
+
+ self.num_layers = num_decoder_layers
+ self.layers = nn.ModuleList(
+ [
+ DecoderLayer(
+ d_model=d_model,
+ attention_dim=attention_dim,
+ num_heads=num_heads,
+ feedforward_dim=feedforward_dim,
+ memory_dim=memory_dim,
+ dropout=dropout,
+ )
+ for _ in range(num_decoder_layers)
+ ]
+ )
+
+ self.output_layer = nn.Linear(d_model, vocab_size)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ memory: Optional[torch.Tensor] = None,
+ memory_lens: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ x: Input tensor of shape (batch, tgt_len).
+ x_lens: A tensor of shape (batch,) containing the number of tokens in `x`
+ before padding.
+ memory:
+ Memory sequence of shape (batch, src_len, memory_dim).
+ memory_lens:
+ A tensor of shape (batch,) containing the number of frames in
+ `memory` before padding.
+
+ Returns:
+ Decoded token logits before softmax (batch, tgt_len, vocab_size)
+ """
+ x = self.embed(x) # (batch, tgt_len, embed_dim)
+ x = self.pos(x) # (batch, tgt_len, embed_dim)
+
+ x = x.permute(1, 0, 2) # (tgt_len, batch, embed_dim)
+
+ # construct attn_mask for self-attn modules
+ padding_mask = make_pad_mask(x_lens) # (batch, tgt_len)
+ causal_mask = subsequent_mask(x.shape[0], device=x.device) # (seq_len, seq_len)
+ attn_mask = torch.logical_or(
+ padding_mask.unsqueeze(1), # (batch, 1, seq_len)
+ torch.logical_not(causal_mask).unsqueeze(0), # (1, seq_len, seq_len)
+ ) # (batch, seq_len, seq_len)
+
+ if memory is not None:
+ memory = memory.permute(1, 0, 2) # (src_len, batch, memory_dim)
+ # construct memory_attn_mask for cross-attn modules
+ memory_padding_mask = make_pad_mask(memory_lens) # (batch, src_len)
+ memory_attn_mask = memory_padding_mask.unsqueeze(1) # (batch, 1, src_len)
+ else:
+ memory_attn_mask = None
+
+ for i, mod in enumerate(self.layers):
+ x = mod(
+ x,
+ attn_mask=attn_mask,
+ memory=memory,
+ memory_attn_mask=memory_attn_mask,
+ )
+
+ x = x.permute(1, 0, 2) # (batch, tgt_len, vocab_size)
+ x = self.output_layer(x)
+
+ return x
+
+
+class DecoderLayer(nn.Module):
+ """Single decoder layer module.
+
+ Args:
+ d_model: equal to decoder_dim, total dimension of the decoder
+ attention_dim: total dimension of multi head attention
+ num_heads: number of attention heads
+ feedforward_dim: hidden dimension of feed_forward module
+ dropout: dropout rate
+ """
+
+ def __init__(
+ self,
+ d_model: int = 512,
+ attention_dim: int = 512,
+ num_heads: int = 8,
+ feedforward_dim: int = 2048,
+ memory_dim: int = 512,
+ dropout: float = 0.1,
+ ):
+ """Construct an DecoderLayer object."""
+ super(DecoderLayer, self).__init__()
+
+ self.norm_self_attn = nn.LayerNorm(d_model)
+ self.self_attn = MultiHeadAttention(
+ d_model, attention_dim, num_heads, dropout=0.0
+ )
+
+ self.norm_src_attn = nn.LayerNorm(d_model)
+ self.src_attn = MultiHeadAttention(
+ d_model, attention_dim, num_heads, memory_dim=memory_dim, dropout=0.0
+ )
+
+ self.norm_ff = nn.LayerNorm(d_model)
+ self.feed_forward = nn.Sequential(
+ nn.Linear(d_model, feedforward_dim),
+ Swish(),
+ nn.Dropout(dropout),
+ nn.Linear(feedforward_dim, d_model),
+ )
+
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ attn_mask: Optional[torch.Tensor] = None,
+ memory: Optional[torch.Tensor] = None,
+ memory_attn_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ x: Input sequence of shape (seq_len, batch, embed_dim).
+ attn_mask: A binary mask for self-attention module indicating which
+ elements will be filled with -inf.
+ Its shape is (batch, 1, src_len) or (batch, tgt_len, src_len).
+ memory: Memory sequence of shape (seq_len, batch, memory_dim).
+ memory_attn_mask: A binary mask for cross-attention module indicating which
+ elements will be filled with -inf.
+ Its shape is (batch, 1, src_len) or (batch, tgt_len, src_len).
+ """
+ # self-attn module
+ qkv = self.norm_self_attn(x)
+ self_attn_out = self.self_attn(
+ query=qkv, key=qkv, value=qkv, attn_mask=attn_mask
+ )
+ x = x + self.dropout(self_attn_out)
+
+ # cross-attn module
+ q = self.norm_src_attn(x)
+ src_attn_out = self.src_attn(
+ query=q, key=memory, value=memory, attn_mask=memory_attn_mask
+ )
+ x = x + self.dropout(src_attn_out)
+
+ # feed-forward module
+ x = x + self.dropout(self.feed_forward(self.norm_ff(x)))
+
+ return x
+
+
+class MultiHeadAttention(nn.Module):
+ """Multi-Head Attention layer.
+
+ Args:
+ embed_dim: total dimension of the model.
+ attention_dim: dimension in the attention module, but must be a multiple of num_heads.
+ num_heads: number of parallel attention heads.
+ memory_dim: dimension of memory embedding, optional.
+ dropout: a Dropout layer on attn_output_weights.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ attention_dim: int,
+ num_heads: int,
+ memory_dim: Optional[int] = None,
+ dropout: float = 0.0,
+ ):
+ super(MultiHeadAttention, self).__init__()
+ self.embed_dim = embed_dim
+ self.attention_dim = attention_dim
+ self.num_heads = num_heads
+ self.head_dim = attention_dim // num_heads
+ assert self.head_dim * num_heads == attention_dim, (
+ self.head_dim,
+ num_heads,
+ attention_dim,
+ )
+ self.dropout = dropout
+ self.name = None # will be overwritten in training code; for diagnostics.
+
+ self.linear_q = nn.Linear(embed_dim, attention_dim, bias=True)
+ self.linear_k = nn.Linear(
+ embed_dim if memory_dim is None else memory_dim, attention_dim, bias=True
+ )
+ self.linear_v = nn.Linear(
+ embed_dim if memory_dim is None else memory_dim, attention_dim, bias=True
+ )
+
+ self.out_proj = nn.Linear(attention_dim, embed_dim, bias=True)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ key_padding_mask: Optional[torch.Tensor] = None,
+ attn_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ """Compute dot product attention.
+
+ Args:
+ query: Query tensor of shape (tgt_len, batch, embed_dim).
+ key: Key tensor of shape (src_len, batch, embed_dim or memory_dim).
+ value: Value tensor of shape (src_len, batch, embed_dim or memory_dim).
+ key_padding_mask: A binary mask indicating which elements are padding.
+ Its shape is (batch, src_len).
+ attn_mask: A binary mask indicating which elements will be filled with -inf.
+ Its shape is (batch, 1, src_len) or (batch, tgt_len, src_len).
+
+ Returns:
+ Output tensor of shape (tgt_len, batch, embed_dim).
+ """
+ num_heads = self.num_heads
+ head_dim = self.head_dim
+
+ tgt_len, batch, _ = query.shape
+ src_len = key.shape[0]
+
+ q = self.linear_q(query) # (tgt_len, batch, num_heads * head_dim)
+ k = self.linear_k(key) # (src_len, batch, num_heads * head_dim)
+ v = self.linear_v(value) # (src_len, batch, num_heads * head_dim)
+
+ q = q.reshape(tgt_len, batch, num_heads, head_dim)
+ q = q.permute(1, 2, 0, 3) # (batch, head, tgt_len, head_dim)
+ k = k.reshape(src_len, batch, num_heads, head_dim)
+ k = k.permute(1, 2, 3, 0) # (batch, head, head_dim, src_len)
+ v = v.reshape(src_len, batch, num_heads, head_dim)
+ v = v.reshape(src_len, batch * num_heads, head_dim).transpose(0, 1)
+
+ # Note: could remove the scaling operation when using ScaledAdam
+ # (batch, head, tgt_len, src_len)
+ attn_weights = torch.matmul(q, k) / math.sqrt(head_dim)
+
+ # From zipformer.py:
+ # This is a harder way of limiting the attention scores to not be too large.
+ # It incurs a penalty if any of them has an absolute value greater than 50.0.
+ # this should be outside the normal range of the attention scores. We use
+ # this mechanism instead of, say, a limit on entropy, because once the entropy
+ # gets very small gradients through the softmax can become very small, and
+ # some mechanisms like that become ineffective.
+ attn_weights = penalize_abs_values_gt(attn_weights, limit=50.0, penalty=1.0e-04)
+
+ if key_padding_mask is not None:
+ assert key_padding_mask.shape == (batch, src_len), key_padding_mask.shape
+ attn_weights = attn_weights.masked_fill(
+ key_padding_mask.unsqueeze(1).unsqueeze(2),
+ float("-inf"),
+ )
+
+ if attn_mask is not None:
+ assert attn_mask.shape == (batch, 1, src_len) or attn_mask.shape == (
+ batch,
+ tgt_len,
+ src_len,
+ ), attn_mask.shape
+ attn_weights = attn_weights.masked_fill(
+ attn_mask.unsqueeze(1), float("-inf")
+ )
+
+ attn_weights = attn_weights.view(batch * num_heads, tgt_len, src_len)
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+ attn_weights = nn.functional.dropout(
+ attn_weights, p=self.dropout, training=self.training
+ )
+
+ # (batch * head, tgt_len, head_dim)
+ attn_output = torch.bmm(attn_weights, v)
+ assert attn_output.shape == (
+ batch * num_heads,
+ tgt_len,
+ head_dim,
+ ), attn_output.shape
+
+ attn_output = attn_output.transpose(0, 1).contiguous()
+ attn_output = attn_output.view(tgt_len, batch, num_heads * head_dim)
+
+ # (batch, tgt_len, embed_dim)
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output
+
+
+class PositionalEncoding(nn.Module):
+ """Positional encoding.
+ Copied from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/embedding.py#L35.
+
+ Args:
+ d_model (int): Embedding dimension.
+ dropout_rate (float): Dropout rate.
+ max_len (int): Maximum input length.
+ """
+
+ def __init__(self, d_model, dropout_rate, max_len=5000):
+ """Construct an PositionalEncoding object."""
+ super(PositionalEncoding, self).__init__()
+ self.d_model = d_model
+ self.xscale = math.sqrt(self.d_model)
+ self.dropout = torch.nn.Dropout(p=dropout_rate)
+ self.pe = None
+ self.extend_pe(torch.tensor(0.0).expand(1, max_len))
+
+ def extend_pe(self, x):
+ """Reset the positional encodings."""
+ if self.pe is not None:
+ if self.pe.size(1) >= x.size(1):
+ if self.pe.dtype != x.dtype or self.pe.device != x.device:
+ self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+ return
+ pe = torch.zeros(x.size(1), self.d_model)
+ position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
+ div_term = torch.exp(
+ torch.arange(0, self.d_model, 2, dtype=torch.float32)
+ * -(math.log(10000.0) / self.d_model)
+ )
+ pe[:, 0::2] = torch.sin(position * div_term)
+ pe[:, 1::2] = torch.cos(position * div_term)
+ pe = pe.unsqueeze(0)
+ self.pe = pe.to(device=x.device, dtype=x.dtype)
+
+ def forward(self, x: torch.Tensor):
+ """Add positional encoding.
+
+ Args:
+ x (torch.Tensor): Input tensor (batch, time, `*`).
+
+ Returns:
+ torch.Tensor: Encoded tensor (batch, time, `*`).
+ """
+ self.extend_pe(x)
+ x = x * self.xscale + self.pe[:, : x.size(1)]
+ return self.dropout(x)
+
+
+class Swish(torch.nn.Module):
+ """Construct an Swish object."""
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Return Swich activation function."""
+ return x * torch.sigmoid(x)
+
+
+def subsequent_mask(size, device="cpu", dtype=torch.bool):
+ """Create mask for subsequent steps (size, size).
+
+ :param int size: size of mask
+ :param str device: "cpu" or "cuda" or torch.Tensor.device
+ :param torch.dtype dtype: result dtype
+ :rtype: torch.Tensor
+ >>> subsequent_mask(3)
+ [[1, 0, 0],
+ [1, 1, 0],
+ [1, 1, 1]]
+ """
+ ret = torch.ones(size, size, device=device, dtype=dtype)
+ return torch.tril(ret, out=ret)
+
+
+def _test_attention_decoder_model():
+ m = AttentionDecoderModel(
+ vocab_size=500,
+ decoder_dim=512,
+ num_decoder_layers=6,
+ attention_dim=512,
+ num_heads=8,
+ feedforward_dim=2048,
+ memory_dim=384,
+ dropout=0.1,
+ sos_id=1,
+ eos_id=1,
+ ignore_id=-1,
+ )
+
+ num_param = sum([p.numel() for p in m.parameters()])
+ print(f"Number of model parameters: {num_param}")
+
+ m.eval()
+ encoder_out = torch.randn(2, 50, 384)
+ encoder_out_lens = torch.full((2,), 50)
+ token_ids = [[1, 2, 3, 4], [2, 3, 10]]
+
+ nll = m.nll(encoder_out, encoder_out_lens, token_ids)
+ print(nll)
+
+
+if __name__ == "__main__":
+ _test_attention_decoder_model()
diff --git a/egs/europarl_st/SRT/lcma_srt/beam_search.py b/egs/europarl_st/SRT/lcma_srt/beam_search.py
new file mode 100644
index 0000000000..7f12942793
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/beam_search.py
@@ -0,0 +1,3426 @@
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang
+# Xiaoyu Yang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import warnings
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Union
+
+import k2
+import sentencepiece as spm
+import torch
+from torch import nn
+
+from icefall import ContextGraph, ContextState, NgramLm, NgramLmStateCost
+from icefall.decode import Nbest, one_best_decoding
+from icefall.lm_wrapper import LmScorer
+from icefall.rnn_lm.model import RnnLmModel
+from icefall.transformer_lm.model import TransformerLM
+from icefall.utils import (
+ DecodingResults,
+ KeywordResult,
+ add_eos,
+ add_sos,
+ get_texts,
+ get_texts_with_timestamp,
+)
+
+
+def fast_beam_search_one_best(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ temperature: float = 1.0,
+ ilme_scale: float = 0.0,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+ allow_partial: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """It limits the maximum number of symbols per frame to 1.
+
+ A lattice is first obtained using fast beam search, and then
+ the shortest path within the lattice is used as the final output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ temperature=temperature,
+ ilme_scale=ilme_scale,
+ allow_partial=allow_partial,
+ blank_penalty=blank_penalty,
+ )
+
+ best_path = one_best_decoding(lattice)
+
+ if not return_timestamps:
+ return get_texts(best_path)
+ else:
+ return get_texts_with_timestamp(best_path)
+
+
+def fast_beam_search_nbest_LG(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ num_paths: int,
+ nbest_scale: float = 0.5,
+ use_double_scores: bool = True,
+ temperature: float = 1.0,
+ blank_penalty: float = 0.0,
+ ilme_scale: float = 0.0,
+ return_timestamps: bool = False,
+ allow_partial: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """It limits the maximum number of symbols per frame to 1.
+
+ The process to get the results is:
+ - (1) Use fast beam search to get a lattice
+ - (2) Select `num_paths` paths from the lattice using k2.random_paths()
+ - (3) Unique the selected paths
+ - (4) Intersect the selected paths with the lattice and compute the
+ shortest path from the intersection result
+ - (5) The path with the largest score is used as the decoding output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ num_paths:
+ Number of paths to extract from the decoded lattice.
+ nbest_scale:
+ It's the scale applied to the lattice.scores. A smaller value
+ yields more unique paths.
+ use_double_scores:
+ True to use double precision for computation. False to use
+ single precision.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ temperature=temperature,
+ allow_partial=allow_partial,
+ blank_penalty=blank_penalty,
+ ilme_scale=ilme_scale,
+ )
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+
+ # The following code is modified from nbest.intersect()
+ word_fsa = k2.invert(nbest.fsa)
+ if hasattr(lattice, "aux_labels"):
+ # delete token IDs as it is not needed
+ del word_fsa.aux_labels
+ word_fsa.scores.zero_()
+ word_fsa_with_epsilon_loops = k2.linear_fsa_with_self_loops(word_fsa)
+ path_to_utt_map = nbest.shape.row_ids(1)
+
+ if hasattr(lattice, "aux_labels"):
+ # lattice has token IDs as labels and word IDs as aux_labels.
+ # inv_lattice has word IDs as labels and token IDs as aux_labels
+ inv_lattice = k2.invert(lattice)
+ inv_lattice = k2.arc_sort(inv_lattice)
+ else:
+ inv_lattice = k2.arc_sort(lattice)
+
+ if inv_lattice.shape[0] == 1:
+ path_lattice = k2.intersect_device(
+ inv_lattice,
+ word_fsa_with_epsilon_loops,
+ b_to_a_map=torch.zeros_like(path_to_utt_map),
+ sorted_match_a=True,
+ )
+ else:
+ path_lattice = k2.intersect_device(
+ inv_lattice,
+ word_fsa_with_epsilon_loops,
+ b_to_a_map=path_to_utt_map,
+ sorted_match_a=True,
+ )
+
+ # path_lattice has word IDs as labels and token IDs as aux_labels
+ path_lattice = k2.top_sort(k2.connect(path_lattice))
+ tot_scores = path_lattice.get_tot_scores(
+ use_double_scores=use_double_scores,
+ log_semiring=True, # Note: we always use True
+ )
+ # See https://github.com/k2-fsa/icefall/pull/420 for why
+ # we always use log_semiring=True
+
+ ragged_tot_scores = k2.RaggedTensor(nbest.shape, tot_scores)
+ best_hyp_indexes = ragged_tot_scores.argmax()
+ best_path = k2.index_fsa(nbest.fsa, best_hyp_indexes)
+
+ if not return_timestamps:
+ return get_texts(best_path)
+ else:
+ return get_texts_with_timestamp(best_path)
+
+
+def fast_beam_search_nbest(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ num_paths: int,
+ nbest_scale: float = 0.5,
+ use_double_scores: bool = True,
+ temperature: float = 1.0,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+ allow_partial: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """It limits the maximum number of symbols per frame to 1.
+
+ The process to get the results is:
+ - (1) Use fast beam search to get a lattice
+ - (2) Select `num_paths` paths from the lattice using k2.random_paths()
+ - (3) Unique the selected paths
+ - (4) Intersect the selected paths with the lattice and compute the
+ shortest path from the intersection result
+ - (5) The path with the largest score is used as the decoding output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ num_paths:
+ Number of paths to extract from the decoded lattice.
+ nbest_scale:
+ It's the scale applied to the lattice.scores. A smaller value
+ yields more unique paths.
+ use_double_scores:
+ True to use double precision for computation. False to use
+ single precision.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ blank_penalty=blank_penalty,
+ temperature=temperature,
+ allow_partial=allow_partial,
+ )
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+
+ # at this point, nbest.fsa.scores are all zeros.
+
+ nbest = nbest.intersect(lattice)
+ # Now nbest.fsa.scores contains acoustic scores
+
+ max_indexes = nbest.tot_scores().argmax()
+
+ best_path = k2.index_fsa(nbest.fsa, max_indexes)
+
+ if not return_timestamps:
+ return get_texts(best_path)
+ else:
+ return get_texts_with_timestamp(best_path)
+
+
+def fast_beam_search_nbest_oracle(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ num_paths: int,
+ ref_texts: List[List[int]],
+ use_double_scores: bool = True,
+ nbest_scale: float = 0.5,
+ temperature: float = 1.0,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+ allow_partial: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """It limits the maximum number of symbols per frame to 1.
+
+ A lattice is first obtained using fast beam search, and then
+ we select `num_paths` linear paths from the lattice. The path
+ that has the minimum edit distance with the given reference transcript
+ is used as the output.
+
+ This is the best result we can achieve for any nbest based rescoring
+ methods.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ num_paths:
+ Number of paths to extract from the decoded lattice.
+ ref_texts:
+ A list-of-list of integers containing the reference transcripts.
+ If the decoding_graph is a trivial_graph, the integer ID is the
+ BPE token ID.
+ use_double_scores:
+ True to use double precision for computation. False to use
+ single precision.
+ nbest_scale:
+ It's the scale applied to the lattice.scores. A smaller value
+ yields more unique paths.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ temperature=temperature,
+ allow_partial=allow_partial,
+ blank_penalty=blank_penalty,
+ )
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+
+ hyps = nbest.build_levenshtein_graphs()
+ refs = k2.levenshtein_graph(ref_texts, device=hyps.device)
+
+ levenshtein_alignment = k2.levenshtein_alignment(
+ refs=refs,
+ hyps=hyps,
+ hyp_to_ref_map=nbest.shape.row_ids(1),
+ sorted_match_ref=True,
+ )
+
+ tot_scores = levenshtein_alignment.get_tot_scores(
+ use_double_scores=False, log_semiring=False
+ )
+ ragged_tot_scores = k2.RaggedTensor(nbest.shape, tot_scores)
+
+ max_indexes = ragged_tot_scores.argmax()
+
+ best_path = k2.index_fsa(nbest.fsa, max_indexes)
+
+ if not return_timestamps:
+ return get_texts(best_path)
+ else:
+ return get_texts_with_timestamp(best_path)
+
+
+def fast_beam_search(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ temperature: float = 1.0,
+ subtract_ilme: bool = False,
+ ilme_scale: float = 0.1,
+ allow_partial: bool = False,
+ blank_penalty: float = 0.0,
+) -> k2.Fsa:
+ """It limits the maximum number of symbols per frame to 1.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi..
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ temperature:
+ Softmax temperature.
+ Returns:
+ Return an FsaVec with axes [utt][state][arc] containing the decoded
+ lattice. Note: When the input graph is a TrivialGraph, the returned
+ lattice is actually an acceptor.
+ """
+ assert encoder_out.ndim == 3
+
+ context_size = model.decoder.context_size
+ vocab_size = model.decoder.vocab_size
+
+ B, T, C = encoder_out.shape
+
+ config = k2.RnntDecodingConfig(
+ vocab_size=vocab_size,
+ decoder_history_len=context_size,
+ beam=beam,
+ max_contexts=max_contexts,
+ max_states=max_states,
+ )
+ individual_streams = []
+ for i in range(B):
+ individual_streams.append(k2.RnntDecodingStream(decoding_graph))
+ decoding_streams = k2.RnntDecodingStreams(individual_streams, config)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ # shape is a RaggedShape of shape (B, context)
+ # contexts is a Tensor of shape (shape.NumElements(), context_size)
+ shape, contexts = decoding_streams.get_contexts()
+ # `nn.Embedding()` in torch below v1.7.1 supports only torch.int64
+ contexts = contexts.to(torch.int64)
+ # decoder_out is of shape (shape.NumElements(), 1, decoder_out_dim)
+ decoder_out = model.decoder(contexts, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # current_encoder_out is of shape
+ # (shape.NumElements(), 1, joiner_dim)
+ # fmt: off
+ current_encoder_out = torch.index_select(
+ encoder_out[:, t:t + 1, :], 0, shape.row_ids(1).to(torch.int64)
+ )
+ # fmt: on
+ logits = model.joiner(
+ current_encoder_out.unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ logits = logits.squeeze(1).squeeze(1)
+
+ if blank_penalty != 0:
+ logits[:, 0] -= blank_penalty
+
+ log_probs = (logits / temperature).log_softmax(dim=-1)
+
+ if ilme_scale != 0:
+ ilme_logits = model.joiner(
+ torch.zeros_like(
+ current_encoder_out, device=current_encoder_out.device
+ ).unsqueeze(2),
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+ ilme_logits = ilme_logits.squeeze(1).squeeze(1)
+ if blank_penalty != 0:
+ ilme_logits[:, 0] -= blank_penalty
+ ilme_log_probs = (ilme_logits / temperature).log_softmax(dim=-1)
+ log_probs -= ilme_scale * ilme_log_probs
+
+ decoding_streams.advance(log_probs)
+ decoding_streams.terminate_and_flush_to_streams()
+ lattice = decoding_streams.format_output(
+ encoder_out_lens.tolist(), allow_partial=allow_partial
+ )
+
+ return lattice
+
+
+def greedy_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ max_sym_per_frame: int,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+) -> Union[List[int], DecodingResults]:
+ """Greedy search for a single utterance.
+ Args:
+ model:
+ An instance of `Transducer`.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder. Support only N==1 for now.
+ max_sym_per_frame:
+ Maximum number of symbols per frame. If it is set to 0, the WER
+ would be 100%.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ assert encoder_out.ndim == 3
+
+ # support only batch_size == 1 for now
+ assert encoder_out.size(0) == 1, encoder_out.size(0)
+
+ blank_id = model.decoder.blank_id
+ context_size = model.decoder.context_size
+ unk_id = getattr(model, "unk_id", blank_id)
+
+ device = next(model.parameters()).device
+
+ decoder_input = torch.tensor(
+ [-1] * (context_size - 1) + [blank_id], device=device, dtype=torch.int64
+ ).reshape(1, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ T = encoder_out.size(1)
+ t = 0
+ hyp = [blank_id] * context_size
+
+ # timestamp[i] is the frame index after subsampling
+ # on which hyp[i] is decoded
+ timestamp = []
+
+ # Maximum symbols per utterance.
+ max_sym_per_utt = 1000
+
+ # symbols per frame
+ sym_per_frame = 0
+
+ # symbols per utterance decoded so far
+ sym_per_utt = 0
+
+ while t < T and sym_per_utt < max_sym_per_utt:
+ if sym_per_frame >= max_sym_per_frame:
+ sym_per_frame = 0
+ t += 1
+ continue
+
+ # fmt: off
+ current_encoder_out = encoder_out[:, t:t+1, :].unsqueeze(2)
+ # fmt: on
+ logits = model.joiner(
+ current_encoder_out, decoder_out.unsqueeze(1), project_input=False
+ )
+ # logits is (1, 1, 1, vocab_size)
+
+ if blank_penalty != 0:
+ logits[:, :, :, 0] -= blank_penalty
+
+ y = logits.argmax().item()
+ if y not in (blank_id, unk_id):
+ hyp.append(y)
+ timestamp.append(t)
+ decoder_input = torch.tensor([hyp[-context_size:]], device=device).reshape(
+ 1, context_size
+ )
+
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ sym_per_utt += 1
+ sym_per_frame += 1
+ else:
+ sym_per_frame = 0
+ t += 1
+ hyp = hyp[context_size:] # remove blanks
+
+ if not return_timestamps:
+ return hyp
+ else:
+ return DecodingResults(
+ hyps=[hyp],
+ timestamps=[timestamp],
+ )
+
+
+def greedy_search_batch(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ blank_penalty: float = 0,
+ return_timestamps: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """Greedy search in batch mode. It hardcodes --max-sym-per-frame=1.
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C), where N >= 1.
+ encoder_out_lens:
+ A 1-D tensor of shape (N,), containing number of valid frames in
+ encoder_out before padding.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ assert encoder_out.ndim == 3
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ device = next(model.parameters()).device
+
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ hyps = [[-1] * (context_size - 1) + [blank_id] for _ in range(N)]
+
+ # timestamp[n][i] is the frame index after subsampling
+ # on which hyp[n][i] is decoded
+ timestamps = [[] for _ in range(N)]
+ # scores[n][i] is the logits on which hyp[n][i] is decoded
+ scores = [[] for _ in range(N)]
+
+ decoder_input = torch.tensor(
+ hyps,
+ device=device,
+ dtype=torch.int64,
+ ) # (N, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_out: (N, 1, decoder_out_dim)
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape: (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ decoder_out = decoder_out[:batch_size]
+
+ logits = model.joiner(
+ current_encoder_out, decoder_out.unsqueeze(1), project_input=False
+ )
+ # logits'shape (batch_size, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (batch_size, vocab_size)
+ assert logits.ndim == 2, logits.shape
+
+ if blank_penalty != 0:
+ logits[:, 0] -= blank_penalty
+
+ y = logits.argmax(dim=1).tolist()
+ emitted = False
+ for i, v in enumerate(y):
+ if v not in (blank_id, unk_id):
+ hyps[i].append(v)
+ timestamps[i].append(t)
+ scores[i].append(logits[i, v].item())
+ emitted = True
+ if emitted:
+ # update decoder output
+ decoder_input = [h[-context_size:] for h in hyps[:batch_size]]
+ decoder_input = torch.tensor(
+ decoder_input,
+ device=device,
+ dtype=torch.int64,
+ )
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ sorted_ans = [h[context_size:] for h in hyps]
+ ans = []
+ ans_timestamps = []
+ ans_scores = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+ ans_timestamps.append(timestamps[unsorted_indices[i]])
+ ans_scores.append(scores[unsorted_indices[i]])
+
+ if not return_timestamps:
+ return ans
+ else:
+ return DecodingResults(
+ hyps=ans,
+ timestamps=ans_timestamps,
+ scores=ans_scores,
+ )
+
+
+@dataclass
+class Hypothesis:
+ # The predicted tokens so far.
+ # Newly predicted tokens are appended to `ys`.
+ ys: List[int]
+
+ # The log prob of ys.
+ # It contains only one entry.
+ log_prob: torch.Tensor
+
+ ac_probs: Optional[List[float]] = None
+
+ # timestamp[i] is the frame index after subsampling
+ # on which ys[i] is decoded
+ timestamp: List[int] = field(default_factory=list)
+
+ # the lm score for next token given the current ys
+ lm_score: Optional[torch.Tensor] = None
+
+ # the RNNLM states (h and c in LSTM)
+ state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+
+ # N-gram LM state
+ state_cost: Optional[NgramLmStateCost] = None
+
+ # Context graph state
+ context_state: Optional[ContextState] = None
+
+ num_tailing_blanks: int = 0
+
+ @property
+ def key(self) -> str:
+ """Return a string representation of self.ys"""
+ return "_".join(map(str, self.ys))
+
+
+class HypothesisList(object):
+ def __init__(self, data: Optional[Dict[str, Hypothesis]] = None) -> None:
+ """
+ Args:
+ data:
+ A dict of Hypotheses. Its key is its `value.key`.
+ """
+ if data is None:
+ self._data = {}
+ else:
+ self._data = data
+
+ @property
+ def data(self) -> Dict[str, Hypothesis]:
+ return self._data
+
+ def add(self, hyp: Hypothesis) -> None:
+ """Add a Hypothesis to `self`.
+
+ If `hyp` already exists in `self`, its probability is updated using
+ `log-sum-exp` with the existed one.
+
+ Args:
+ hyp:
+ The hypothesis to be added.
+ """
+ key = hyp.key
+ if key in self:
+ old_hyp = self._data[key] # shallow copy
+ torch.logaddexp(old_hyp.log_prob, hyp.log_prob, out=old_hyp.log_prob)
+ else:
+ self._data[key] = hyp
+
+ def get_most_probable(self, length_norm: bool = False) -> Hypothesis:
+ """Get the most probable hypothesis, i.e., the one with
+ the largest `log_prob`.
+
+ Args:
+ length_norm:
+ If True, the `log_prob` of a hypothesis is normalized by the
+ number of tokens in it.
+ Returns:
+ Return the hypothesis that has the largest `log_prob`.
+ """
+ if length_norm:
+ return max(self._data.values(), key=lambda hyp: hyp.log_prob / len(hyp.ys))
+ else:
+ return max(self._data.values(), key=lambda hyp: hyp.log_prob)
+
+ def remove(self, hyp: Hypothesis) -> None:
+ """Remove a given hypothesis.
+
+ Caution:
+ `self` is modified **in-place**.
+
+ Args:
+ hyp:
+ The hypothesis to be removed from `self`.
+ Note: It must be contained in `self`. Otherwise,
+ an exception is raised.
+ """
+ key = hyp.key
+ assert key in self, f"{key} does not exist"
+ del self._data[key]
+
+ def filter(self, threshold: torch.Tensor) -> "HypothesisList":
+ """Remove all Hypotheses whose log_prob is less than threshold.
+
+ Caution:
+ `self` is not modified. Instead, a new HypothesisList is returned.
+
+ Returns:
+ Return a new HypothesisList containing all hypotheses from `self`
+ with `log_prob` being greater than the given `threshold`.
+ """
+ ans = HypothesisList()
+ for _, hyp in self._data.items():
+ if hyp.log_prob > threshold:
+ ans.add(hyp) # shallow copy
+ return ans
+
+ def topk(self, k: int, length_norm: bool = False) -> "HypothesisList":
+ """Return the top-k hypothesis.
+
+ Args:
+ length_norm:
+ If True, the `log_prob` of a hypothesis is normalized by the
+ number of tokens in it.
+ """
+ hyps = list(self._data.items())
+
+ if length_norm:
+ hyps = sorted(
+ hyps, key=lambda h: h[1].log_prob / len(h[1].ys), reverse=True
+ )[:k]
+ else:
+ hyps = sorted(hyps, key=lambda h: h[1].log_prob, reverse=True)[:k]
+
+ ans = HypothesisList(dict(hyps))
+ return ans
+
+ def __contains__(self, key: str):
+ return key in self._data
+
+ def __iter__(self):
+ return iter(self._data.values())
+
+ def __len__(self) -> int:
+ return len(self._data)
+
+ def __str__(self) -> str:
+ s = []
+ for key in self:
+ s.append(key)
+ return ", ".join(s)
+
+
+def get_hyps_shape(hyps: List[HypothesisList]) -> k2.RaggedShape:
+ """Return a ragged shape with axes [utt][num_hyps].
+
+ Args:
+ hyps:
+ len(hyps) == batch_size. It contains the current hypothesis for
+ each utterance in the batch.
+ Returns:
+ Return a ragged shape with 2 axes [utt][num_hyps]. Note that
+ the shape is on CPU.
+ """
+ num_hyps = [len(h) for h in hyps]
+
+ # torch.cumsum() is inclusive sum, so we put a 0 at the beginning
+ # to get exclusive sum later.
+ num_hyps.insert(0, 0)
+
+ num_hyps = torch.tensor(num_hyps)
+ row_splits = torch.cumsum(num_hyps, dim=0, dtype=torch.int32)
+ ans = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=row_splits[-1].item()
+ )
+ return ans
+
+
+def keywords_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ keywords_graph: ContextGraph,
+ beam: int = 4,
+ num_tailing_blanks: int = 0,
+ blank_penalty: float = 0,
+) -> List[List[KeywordResult]]:
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C).
+ encoder_out_lens:
+ A 1-D tensor of shape (N,), containing number of valid frames in
+ encoder_out before padding.
+ keywords_graph:
+ A instance of ContextGraph containing keywords and their configurations.
+ beam:
+ Number of active paths during the beam search.
+ num_tailing_blanks:
+ The number of tailing blanks a keyword should be followed, this is for the
+ scenario that a keyword will be the prefix of another. In most cases, you
+ can just set it to 0.
+ blank_penalty:
+ The score used to penalize blank probability.
+ Returns:
+ Return a list of list of KeywordResult.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+ assert keywords_graph is not None
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ context_state=keywords_graph.root,
+ timestamp=[],
+ ac_probs=[],
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ sorted_ans = [[] for _ in range(N)]
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ ) # (num_hyps, 1)
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_out is of shape (num_hyps, 1, 1, joiner_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ if blank_penalty != 0:
+ logits[:, 0] -= blank_penalty
+
+ probs = logits.softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs = probs.log()
+
+ probs = probs.reshape(-1)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+ ragged_probs = k2.RaggedTensor(shape=log_probs_shape, value=probs)
+
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+ hyp_probs = ragged_probs[i].tolist()
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+ new_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ new_timestamp = hyp.timestamp[:]
+ new_ac_probs = hyp.ac_probs[:]
+ context_score = 0
+ new_context_state = hyp.context_state
+ new_num_tailing_blanks = hyp.num_tailing_blanks + 1
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ new_timestamp.append(t)
+ new_ac_probs.append(hyp_probs[topk_indexes[k]])
+ (
+ context_score,
+ new_context_state,
+ _,
+ ) = keywords_graph.forward_one_step(hyp.context_state, new_token)
+ new_num_tailing_blanks = 0
+ if new_context_state.token == -1: # root
+ new_ys[-context_size:] = [-1] * (context_size - 1) + [blank_id]
+
+ new_log_prob = topk_log_probs[k] + context_score
+
+ new_hyp = Hypothesis(
+ ys=new_ys,
+ log_prob=new_log_prob,
+ timestamp=new_timestamp,
+ ac_probs=new_ac_probs,
+ context_state=new_context_state,
+ num_tailing_blanks=new_num_tailing_blanks,
+ )
+ B[i].add(new_hyp)
+
+ top_hyp = B[i].get_most_probable(length_norm=True)
+ matched, matched_state = keywords_graph.is_matched(top_hyp.context_state)
+ if matched:
+ ac_prob = (
+ sum(top_hyp.ac_probs[-matched_state.level :]) / matched_state.level
+ )
+ if (
+ matched
+ and top_hyp.num_tailing_blanks > num_tailing_blanks
+ and ac_prob >= matched_state.ac_threshold
+ ):
+ keyword = KeywordResult(
+ hyps=top_hyp.ys[-matched_state.level :],
+ timestamps=top_hyp.timestamp[-matched_state.level :],
+ phrase=matched_state.phrase,
+ )
+ sorted_ans[i].append(keyword)
+ B[i] = HypothesisList()
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ context_state=keywords_graph.root,
+ timestamp=[],
+ ac_probs=[],
+ )
+ )
+
+ B = B + finalized_B
+
+ for i, hyps in enumerate(B):
+ top_hyp = hyps.get_most_probable(length_norm=True)
+ matched, matched_state = keywords_graph.is_matched(top_hyp.context_state)
+ if matched:
+ ac_prob = (
+ sum(top_hyp.ac_probs[-matched_state.level :]) / matched_state.level
+ )
+ if matched and ac_prob >= matched_state.ac_threshold:
+ keyword = KeywordResult(
+ hyps=top_hyp.ys[-matched_state.level :],
+ timestamps=top_hyp.timestamp[-matched_state.level :],
+ phrase=matched_state.phrase,
+ )
+ sorted_ans[i].append(keyword)
+
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+ return ans
+
+
+# def modified_beam_search(
+# model: nn.Module,
+# encoder_out: torch.Tensor,
+# encoder_out_lens: torch.Tensor,
+# decoder: Optional[nn.Module] = None,
+# joiner: Optional[nn.Module] = None,
+# context_graph: Optional[ContextGraph] = None,
+# beam: int = 4,
+# temperature: float = 1.0,
+# blank_penalty: float = 0.0,
+# return_timestamps: bool = False,
+# ) -> Union[List[List[int]], DecodingResults]:
+# """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+# Args:
+# model:
+# The transducer model.
+# encoder_out:
+# Output from the encoder. Its shape is (N, T, C).
+# encoder_out_lens:
+# A 1-D tensor of shape (N,), containing number of valid frames in
+# encoder_out before padding.
+# beam:
+# Number of active paths during the beam search.
+# temperature:
+# Softmax temperature.
+# return_timestamps:
+# Whether to return timestamps.
+# Returns:
+# If return_timestamps is False, return the decoded result.
+# Else, return a DecodingResults object containing
+# decoded result and corresponding timestamps.
+# """
+# assert encoder_out.ndim == 3, encoder_out.shape
+# assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+# packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+# input=encoder_out,
+# lengths=encoder_out_lens.cpu(),
+# batch_first=True,
+# enforce_sorted=False,
+# )
+
+# blank_id = decoder.blank_id
+# unk_id = getattr(model, "unk_id", blank_id)
+# context_size = decoder.context_size
+# device = next(model.parameters()).device
+
+# batch_size_list = packed_encoder_out.batch_sizes.tolist()
+# N = encoder_out.size(0)
+# assert torch.all(encoder_out_lens > 0), encoder_out_lens
+# assert N == batch_size_list[0], (N, batch_size_list)
+
+# B = [HypothesisList() for _ in range(N)]
+# for i in range(N):
+# B[i].add(
+# Hypothesis(
+# ys=[-1] * (context_size - 1) + [blank_id],
+# log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+# context_state=None if context_graph is None else context_graph.root,
+# timestamp=[],
+# )
+# )
+
+# encoder_out = joiner.encoder_proj(packed_encoder_out.data)
+
+# offset = 0
+# finalized_B = []
+# for t, batch_size in enumerate(batch_size_list):
+# start = offset
+# end = offset + batch_size
+# current_encoder_out = encoder_out.data[start:end]
+# current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+# # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+# offset = end
+
+# finalized_B = B[batch_size:] + finalized_B
+# B = B[:batch_size]
+
+# hyps_shape = get_hyps_shape(B).to(device)
+
+# A = [list(b) for b in B]
+
+# B = [HypothesisList() for _ in range(batch_size)]
+
+# ys_log_probs = torch.cat(
+# [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+# ) # (num_hyps, 1)
+
+# decoder_input = torch.tensor(
+# [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+# device=device,
+# dtype=torch.int64,
+# ) # (num_hyps, context_size)
+
+# decoder_out = decoder(decoder_input, need_pad=False).unsqueeze(1)
+# decoder_out = joiner.decoder_proj(decoder_out)
+# # decoder_out is of shape (num_hyps, 1, 1, joiner_dim)
+
+# # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+# # as index, so we use `to(torch.int64)` below.
+# current_encoder_out = torch.index_select(
+# current_encoder_out,
+# dim=0,
+# index=hyps_shape.row_ids(1).to(torch.int64),
+# ) # (num_hyps, 1, 1, encoder_out_dim)
+
+# logits = joiner(
+# current_encoder_out,
+# decoder_out,
+# project_input=False,
+# ) # (num_hyps, 1, 1, vocab_size)
+
+# logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+# if blank_penalty != 0:
+# logits[:, 0] -= blank_penalty
+
+# log_probs = (logits / temperature).log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+# log_probs.add_(ys_log_probs)
+
+# vocab_size = log_probs.size(-1)
+
+# log_probs = log_probs.reshape(-1)
+
+# row_splits = hyps_shape.row_splits(1) * vocab_size
+# log_probs_shape = k2.ragged.create_ragged_shape2(
+# row_splits=row_splits, cached_tot_size=log_probs.numel()
+# )
+# ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+
+# for i in range(batch_size):
+# topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+# with warnings.catch_warnings():
+# warnings.simplefilter("ignore")
+# topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+# topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+# for k in range(len(topk_hyp_indexes)):
+# hyp_idx = topk_hyp_indexes[k]
+# hyp = A[i][hyp_idx]
+# new_ys = hyp.ys[:]
+# new_token = topk_token_indexes[k]
+# new_timestamp = hyp.timestamp[:]
+# context_score = 0
+# new_context_state = None if context_graph is None else hyp.context_state
+# if new_token not in (blank_id, unk_id):
+# new_ys.append(new_token)
+# new_timestamp.append(t)
+# if context_graph is not None:
+# (
+# context_score,
+# new_context_state,
+# ) = context_graph.forward_one_step(hyp.context_state, new_token)
+
+# new_log_prob = topk_log_probs[k] + context_score
+
+# new_hyp = Hypothesis(
+# ys=new_ys,
+# log_prob=new_log_prob,
+# timestamp=new_timestamp,
+# context_state=new_context_state,
+# )
+# B[i].add(new_hyp)
+
+# B = B + finalized_B
+
+# # finalize context_state, if the matched contexts do not reach final state
+# # we need to add the score on the corresponding backoff arc
+# if context_graph is not None:
+# finalized_B = [HypothesisList() for _ in range(len(B))]
+# for i, hyps in enumerate(B):
+# for hyp in list(hyps):
+# context_score, new_context_state = context_graph.finalize(
+# hyp.context_state
+# )
+# finalized_B[i].add(
+# Hypothesis(
+# ys=hyp.ys,
+# log_prob=hyp.log_prob + context_score,
+# timestamp=hyp.timestamp,
+# context_state=new_context_state,
+# )
+# )
+# B = finalized_B
+
+# best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+# sorted_ans = [h.ys[context_size:] for h in best_hyps]
+# sorted_timestamps = [h.timestamp for h in best_hyps]
+# ans = []
+# ans_timestamps = []
+# unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+# for i in range(N):
+# ans.append(sorted_ans[unsorted_indices[i]])
+# ans_timestamps.append(sorted_timestamps[unsorted_indices[i]])
+
+# if not return_timestamps:
+# return ans
+# else:
+# return DecodingResults(
+# hyps=ans,
+# timestamps=ans_timestamps,
+# )
+
+
+def modified_beam_search(
+ model: torch.nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ decoder: Optional[torch.nn.Module] = None,
+ joiner: Optional[torch.nn.Module] = None,
+ context_graph: Optional["ContextGraph"] = None,
+ beam: int = 4,
+ temperature: float = 1.0,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+ # ===== Additional parameters (same target lang for entire batch) =====
+ lang_token_id: Optional[int] = None,
+ force_first_lang: bool = False,
+) -> Union[List[List[int]], "DecodingResults"]:
+ """
+ Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. (N, T, C)
+ encoder_out_lens:
+ Valid frame lengths. (N,)
+ beam:
+ Beam width.
+ temperature:
+ Softmax temperature.
+ blank_penalty:
+ Logit penalty applied to blank.
+ return_timestamps:
+ Whether to return token timestamps.
+ lang_token_id:
+ Target language token id (same for entire batch), e.g. sp_st.piece_to_id("<2zh-cn>").
+ force_first_lang:
+ If True: before emitting the first label, only allow {blank, lang_token_id}.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+
+ if force_first_lang:
+ if lang_token_id is None:
+ raise ValueError("force_first_lang=True but lang_token_id not provided")
+ if not isinstance(lang_token_id, int):
+ raise TypeError("lang_token_id must be int (same target lang for entire batch).")
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ # Initialize beam for each utterance (use blank as SOS, fill context_size)
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ context_state=None if context_graph is None else context_graph.root,
+ timestamp=[],
+ )
+ )
+
+ # Pre-compute encoder projection for speedup
+ enc_proj = joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = enc_proj[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # (batch_size, 1, 1, joiner_dim)
+ offset = end
+
+ # Store beams of finished samples
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+ A = [list(b) for b in B] # Current hypotheses for all samples
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ # Accumulated scores (all hypotheses stacked)
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ ) # (num_hyps, 1)
+
+ # Take last context_size tokens of each hypothesis as predictor input
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = joiner.decoder_proj(decoder_out) # (num_hyps, 1, 1, joiner_dim)
+
+ # Expand current frame encoder representation to each hypothesis
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, joiner_dim)
+
+ # Pass through joiner to get vocab logits
+ logits = joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, V)
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, V)
+
+ # ====== First-emit constraint: only allow {blank, lang_token_id} before first emission ======
+ if force_first_lang:
+ # When len(hyp.ys) == context_size, no label has been emitted yet (except initial context)
+ need_force = torch.tensor(
+ [len(hyp.ys) == context_size for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.bool,
+ ) # (num_hyps,)
+ if need_force.any():
+ vocab_size = logits.size(-1)
+ if not (0 <= lang_token_id < vocab_size):
+ raise RuntimeError("lang_token_id out of vocabulary range")
+ logits_backup = logits.clone()
+ logits[need_force] = float("-inf")
+ rows = torch.nonzero(need_force, as_tuple=True)[0]
+ # Only allow raw scores of blank and target lang token
+ logits[rows, blank_id] = logits_backup[rows, blank_id]
+ logits[rows, lang_token_id] = logits_backup[rows, lang_token_id]
+ # ============================================================
+
+ # Optional blank penalty (applied after first-emit constraint, can stack)
+ if blank_penalty != 0:
+ logits[:, blank_id] -= blank_penalty
+
+ # Softmax(temperature) + accumulated scores
+ log_probs = (logits / temperature).log_softmax(dim=-1) # (num_hyps, V)
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+
+ # Select top-k for each sample and expand new hypotheses
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+ new_token = topk_token_indexes[k]
+
+ new_ys = hyp.ys[:]
+ new_timestamp = hyp.timestamp[:]
+ context_score = 0
+ new_context_state = None if context_graph is None else hyp.context_state
+
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ new_timestamp.append(t)
+ if context_graph is not None:
+ (
+ context_score,
+ new_context_state,
+ ) = context_graph.forward_one_step(hyp.context_state, new_token)
+
+ new_log_prob = topk_log_probs[k] + context_score
+ B[i].add(
+ Hypothesis(
+ ys=new_ys,
+ log_prob=new_log_prob,
+ timestamp=new_timestamp,
+ context_state=new_context_state,
+ )
+ )
+
+ B = B + finalized_B
+
+ # Finalize context_graph at the end (add backoff scores)
+ if context_graph is not None:
+ finalized_B = [HypothesisList() for _ in range(len(B))]
+ for i, hyps in enumerate(B):
+ for hyp in list(hyps):
+ context_score, new_context_state = context_graph.finalize(
+ hyp.context_state
+ )
+ finalized_B[i].add(
+ Hypothesis(
+ ys=hyp.ys,
+ log_prob=hyp.log_prob + context_score,
+ timestamp=hyp.timestamp,
+ context_state=new_context_state,
+ )
+ )
+ B = finalized_B
+
+ # Get best path for each utterance (with length normalization)
+ best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+ # Remove predictor context prefix
+ sorted_ans = [h.ys[context_size:] for h in best_hyps]
+ sorted_timestamps = [h.timestamp for h in best_hyps]
+
+ # Restore to original sample order before packing
+ ans, ans_timestamps = [], []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+ ans_timestamps.append(sorted_timestamps[unsorted_indices[i]])
+
+ if not return_timestamps:
+ return ans
+ else:
+ return DecodingResults(hyps=ans, timestamps=ans_timestamps)
+
+
+def modified_beam_search_lm_rescore(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ LM: LmScorer,
+ lm_scale_list: List[int],
+ beam: int = 4,
+ temperature: float = 1.0,
+ return_timestamps: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+ Rescore the final results with RNNLM and return the one with the highest score
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C).
+ encoder_out_lens:
+ A 1-D tensor of shape (N,), containing number of valid frames in
+ encoder_out before padding.
+ beam:
+ Number of active paths during the beam search.
+ temperature:
+ Softmax temperature.
+ LM:
+ A neural network language model
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ timestamp=[],
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ ) # (num_hyps, 1)
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_out is of shape (num_hyps, 1, 1, joiner_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = (logits / temperature).log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ new_timestamp = hyp.timestamp[:]
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ new_timestamp.append(t)
+
+ new_log_prob = topk_log_probs[k]
+ new_hyp = Hypothesis(
+ ys=new_ys, log_prob=new_log_prob, timestamp=new_timestamp
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+
+ # get the am_scores for n-best list
+ hyps_shape = get_hyps_shape(B)
+ am_scores = torch.tensor([hyp.log_prob.item() for b in B for hyp in b])
+ am_scores = k2.RaggedTensor(value=am_scores, shape=hyps_shape).to(device)
+
+ # now LM rescore
+ # prepare input data to LM
+ candidate_seqs = [hyp.ys[context_size:] for b in B for hyp in b]
+ possible_seqs = k2.RaggedTensor(candidate_seqs)
+ row_splits = possible_seqs.shape.row_splits(1)
+ sentence_token_lengths = row_splits[1:] - row_splits[:-1]
+ possible_seqs_with_sos = add_sos(possible_seqs, sos_id=1)
+ possible_seqs_with_eos = add_eos(possible_seqs, eos_id=1)
+ sentence_token_lengths += 1
+
+ x = possible_seqs_with_sos.pad(mode="constant", padding_value=blank_id)
+ y = possible_seqs_with_eos.pad(mode="constant", padding_value=blank_id)
+ x = x.to(device).to(torch.int64)
+ y = y.to(device).to(torch.int64)
+ sentence_token_lengths = sentence_token_lengths.to(device).to(torch.int64)
+
+ lm_scores = LM.lm(x=x, y=y, lengths=sentence_token_lengths)
+ assert lm_scores.ndim == 2
+ lm_scores = -1 * lm_scores.sum(dim=1)
+
+ ans = {}
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+
+ # get the best hyp with different lm_scale
+ for lm_scale in lm_scale_list:
+ key = f"nnlm_scale_{lm_scale:.2f}"
+ tot_scores = am_scores.values + lm_scores * lm_scale
+ ragged_tot_scores = k2.RaggedTensor(shape=am_scores.shape, value=tot_scores)
+ max_indexes = ragged_tot_scores.argmax().tolist()
+ unsorted_hyps = [candidate_seqs[idx] for idx in max_indexes]
+ hyps = []
+ for idx in unsorted_indices:
+ hyps.append(unsorted_hyps[idx])
+
+ ans[key] = hyps
+ return ans
+
+
+def modified_beam_search_lm_rescore_LODR(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ LM: LmScorer,
+ LODR_lm: NgramLm,
+ sp: spm.SentencePieceProcessor,
+ lm_scale_list: List[int],
+ beam: int = 4,
+ temperature: float = 1.0,
+ return_timestamps: bool = False,
+) -> Union[List[List[int]], DecodingResults]:
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+ Rescore the final results with RNNLM and return the one with the highest score
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C).
+ encoder_out_lens:
+ A 1-D tensor of shape (N,), containing number of valid frames in
+ encoder_out before padding.
+ beam:
+ Number of active paths during the beam search.
+ temperature:
+ Softmax temperature.
+ LM:
+ A neural network language model
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ timestamp=[],
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ ) # (num_hyps, 1)
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_out is of shape (num_hyps, 1, 1, joiner_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = (logits / temperature).log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ new_timestamp = hyp.timestamp[:]
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ new_timestamp.append(t)
+
+ new_log_prob = topk_log_probs[k]
+ new_hyp = Hypothesis(
+ ys=new_ys, log_prob=new_log_prob, timestamp=new_timestamp
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+
+ # get the am_scores for n-best list
+ hyps_shape = get_hyps_shape(B)
+ am_scores = torch.tensor([hyp.log_prob.item() for b in B for hyp in b])
+ am_scores = k2.RaggedTensor(value=am_scores, shape=hyps_shape).to(device)
+
+ # now LM rescore
+ # prepare input data to LM
+ candidate_seqs = [hyp.ys[context_size:] for b in B for hyp in b]
+ possible_seqs = k2.RaggedTensor(candidate_seqs)
+ row_splits = possible_seqs.shape.row_splits(1)
+ sentence_token_lengths = row_splits[1:] - row_splits[:-1]
+ possible_seqs_with_sos = add_sos(possible_seqs, sos_id=1)
+ possible_seqs_with_eos = add_eos(possible_seqs, eos_id=1)
+ sentence_token_lengths += 1
+
+ x = possible_seqs_with_sos.pad(mode="constant", padding_value=blank_id)
+ y = possible_seqs_with_eos.pad(mode="constant", padding_value=blank_id)
+ x = x.to(device).to(torch.int64)
+ y = y.to(device).to(torch.int64)
+ sentence_token_lengths = sentence_token_lengths.to(device).to(torch.int64)
+
+ lm_scores = LM.lm(x=x, y=y, lengths=sentence_token_lengths)
+ assert lm_scores.ndim == 2
+ lm_scores = -1 * lm_scores.sum(dim=1)
+
+ # now LODR scores
+ import math
+
+ LODR_scores = []
+ for seq in candidate_seqs:
+ tokens = " ".join(sp.id_to_piece(seq))
+ LODR_scores.append(LODR_lm.score(tokens))
+ LODR_scores = torch.tensor(LODR_scores).to(device) * math.log(
+ 10
+ ) # arpa scores are 10-based
+ assert lm_scores.shape == LODR_scores.shape
+
+ ans = {}
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+
+ LODR_scale_list = [0.05 * i for i in range(1, 20)]
+ # get the best hyp with different lm_scale and lodr_scale
+ for lm_scale in lm_scale_list:
+ for lodr_scale in LODR_scale_list:
+ key = f"nnlm_scale_{lm_scale:.2f}_lodr_scale_{lodr_scale:.2f}"
+ tot_scores = (
+ am_scores.values / lm_scale + lm_scores - LODR_scores * lodr_scale
+ )
+ ragged_tot_scores = k2.RaggedTensor(shape=am_scores.shape, value=tot_scores)
+ max_indexes = ragged_tot_scores.argmax().tolist()
+ unsorted_hyps = [candidate_seqs[idx] for idx in max_indexes]
+ hyps = []
+ for idx in unsorted_indices:
+ hyps.append(unsorted_hyps[idx])
+
+ ans[key] = hyps
+ return ans
+
+
+def _deprecated_modified_beam_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ beam: int = 4,
+ return_timestamps: bool = False,
+) -> Union[List[int], DecodingResults]:
+ """It limits the maximum number of symbols per frame to 1.
+
+ It decodes only one utterance at a time. We keep it only for reference.
+ The function :func:`modified_beam_search` should be preferred as it
+ supports batch decoding.
+
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder. Support only N==1 for now.
+ beam:
+ Beam size.
+ return_timestamps:
+ Whether to return timestamps.
+
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+
+ assert encoder_out.ndim == 3
+
+ # support only batch_size == 1 for now
+ assert encoder_out.size(0) == 1, encoder_out.size(0)
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+
+ device = next(model.parameters()).device
+
+ T = encoder_out.size(1)
+
+ B = HypothesisList()
+ B.add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ timestamp=[],
+ )
+ )
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ for t in range(T):
+ # fmt: off
+ current_encoder_out = encoder_out[:, t:t+1, :].unsqueeze(2)
+ # current_encoder_out is of shape (1, 1, 1, encoder_out_dim)
+ # fmt: on
+ A = list(B)
+ B = HypothesisList()
+
+ ys_log_probs = torch.cat([hyp.log_prob.reshape(1, 1) for hyp in A])
+ # ys_log_probs is of shape (num_hyps, 1)
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyp in A],
+ device=device,
+ dtype=torch.int64,
+ )
+ # decoder_input is of shape (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_output is of shape (num_hyps, 1, 1, joiner_dim)
+
+ current_encoder_out = current_encoder_out.expand(
+ decoder_out.size(0), 1, 1, -1
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ )
+ # logits is of shape (num_hyps, 1, 1, vocab_size)
+ logits = logits.squeeze(1).squeeze(1)
+
+ # now logits is of shape (num_hyps, vocab_size)
+ log_probs = logits.log_softmax(dim=-1)
+
+ log_probs.add_(ys_log_probs)
+
+ log_probs = log_probs.reshape(-1)
+ topk_log_probs, topk_indexes = log_probs.topk(beam)
+
+ # topk_hyp_indexes are indexes into `A`
+ topk_hyp_indexes = topk_indexes // logits.size(-1)
+ topk_token_indexes = topk_indexes % logits.size(-1)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = topk_hyp_indexes.tolist()
+ topk_token_indexes = topk_token_indexes.tolist()
+
+ for i in range(len(topk_hyp_indexes)):
+ hyp = A[topk_hyp_indexes[i]]
+ new_ys = hyp.ys[:]
+ new_timestamp = hyp.timestamp[:]
+ new_token = topk_token_indexes[i]
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ new_timestamp.append(t)
+ new_log_prob = topk_log_probs[i]
+ new_hyp = Hypothesis(
+ ys=new_ys, log_prob=new_log_prob, timestamp=new_timestamp
+ )
+ B.add(new_hyp)
+
+ best_hyp = B.get_most_probable(length_norm=True)
+ ys = best_hyp.ys[context_size:] # [context_size:] to remove blanks
+
+ if not return_timestamps:
+ return ys
+ else:
+ return DecodingResults(hyps=[ys], timestamps=[best_hyp.timestamp])
+
+
+def beam_search(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ beam: int = 4,
+ temperature: float = 1.0,
+ blank_penalty: float = 0.0,
+ return_timestamps: bool = False,
+) -> Union[List[int], DecodingResults]:
+ """
+ It implements Algorithm 1 in https://arxiv.org/pdf/1211.3711.pdf
+
+ espnet/nets/beam_search_transducer.py#L247 is used as a reference.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder. Support only N==1 for now.
+ beam:
+ Beam size.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+
+ Returns:
+ If return_timestamps is False, return the decoded result.
+ Else, return a DecodingResults object containing
+ decoded result and corresponding timestamps.
+ """
+ assert encoder_out.ndim == 3
+
+ # support only batch_size == 1 for now
+ assert encoder_out.size(0) == 1, encoder_out.size(0)
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+
+ device = next(model.parameters()).device
+
+ decoder_input = torch.tensor(
+ [blank_id] * context_size,
+ device=device,
+ dtype=torch.int64,
+ ).reshape(1, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ encoder_out = model.joiner.encoder_proj(encoder_out)
+
+ T = encoder_out.size(1)
+ t = 0
+
+ B = HypothesisList()
+ B.add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id], log_prob=0.0, timestamp=[]
+ )
+ )
+
+ max_sym_per_utt = 20000
+
+ sym_per_utt = 0
+
+ decoder_cache: Dict[str, torch.Tensor] = {}
+
+ while t < T and sym_per_utt < max_sym_per_utt:
+ # fmt: off
+ current_encoder_out = encoder_out[:, t:t+1, :].unsqueeze(2)
+ # fmt: on
+ A = B
+ B = HypothesisList()
+
+ joint_cache: Dict[str, torch.Tensor] = {}
+
+ # TODO(fangjun): Implement prefix search to update the `log_prob`
+ # of hypotheses in A
+
+ while True:
+ y_star = A.get_most_probable()
+ A.remove(y_star)
+
+ cached_key = y_star.key
+
+ if cached_key not in decoder_cache:
+ decoder_input = torch.tensor(
+ [y_star.ys[-context_size:]],
+ device=device,
+ dtype=torch.int64,
+ ).reshape(1, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ decoder_cache[cached_key] = decoder_out
+ else:
+ decoder_out = decoder_cache[cached_key]
+
+ cached_key += f"-t-{t}"
+ if cached_key not in joint_cache:
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out.unsqueeze(1),
+ project_input=False,
+ )
+
+ if blank_penalty != 0:
+ logits[:, :, :, 0] -= blank_penalty
+
+ # TODO(fangjun): Scale the blank posterior
+ log_prob = (logits / temperature).log_softmax(dim=-1)
+ # log_prob is (1, 1, 1, vocab_size)
+ log_prob = log_prob.squeeze()
+ # Now log_prob is (vocab_size,)
+ joint_cache[cached_key] = log_prob
+ else:
+ log_prob = joint_cache[cached_key]
+
+ # First, process the blank symbol
+ skip_log_prob = log_prob[blank_id]
+ new_y_star_log_prob = y_star.log_prob + skip_log_prob
+
+ # ys[:] returns a copy of ys
+ B.add(
+ Hypothesis(
+ ys=y_star.ys[:],
+ log_prob=new_y_star_log_prob,
+ timestamp=y_star.timestamp[:],
+ )
+ )
+
+ # Second, process other non-blank labels
+ values, indices = log_prob.topk(beam + 1)
+ for i, v in zip(indices.tolist(), values.tolist()):
+ if i in (blank_id, unk_id):
+ continue
+ new_ys = y_star.ys + [i]
+ new_log_prob = y_star.log_prob + v
+ new_timestamp = y_star.timestamp + [t]
+ A.add(
+ Hypothesis(
+ ys=new_ys,
+ log_prob=new_log_prob,
+ timestamp=new_timestamp,
+ )
+ )
+
+ # Check whether B contains more than "beam" elements more probable
+ # than the most probable in A
+ A_most_probable = A.get_most_probable()
+
+ kept_B = B.filter(A_most_probable.log_prob)
+
+ if len(kept_B) >= beam:
+ B = kept_B.topk(beam)
+ break
+
+ t += 1
+
+ best_hyp = B.get_most_probable(length_norm=True)
+ ys = best_hyp.ys[context_size:] # [context_size:] to remove blanks
+
+ if not return_timestamps:
+ return ys
+ else:
+ return DecodingResults(hyps=[ys], timestamps=[best_hyp.timestamp])
+
+
+def fast_beam_search_with_nbest_rescoring(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ ngram_lm_scale_list: List[float],
+ num_paths: int,
+ G: k2.Fsa,
+ sp: spm.SentencePieceProcessor,
+ word_table: k2.SymbolTable,
+ oov_word: str = "",
+ use_double_scores: bool = True,
+ nbest_scale: float = 0.5,
+ temperature: float = 1.0,
+ return_timestamps: bool = False,
+) -> Dict[str, Union[List[List[int]], DecodingResults]]:
+ """It limits the maximum number of symbols per frame to 1.
+ A lattice is first obtained using fast beam search, num_path are selected
+ and rescored using a given language model. The shortest path within the
+ lattice is used as the final output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi.
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ ngram_lm_scale_list:
+ A list of floats representing LM score scales.
+ num_paths:
+ Number of paths to extract from the decoded lattice.
+ G:
+ An FsaVec containing only a single FSA. It is an n-gram LM.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ oov_word:
+ OOV words are replaced with this word.
+ use_double_scores:
+ True to use double precision for computation. False to use
+ single precision.
+ nbest_scale:
+ It's the scale applied to the lattice.scores. A smaller value
+ yields more unique paths.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ Return the decoded result in a dict, where the key has the form
+ 'ngram_lm_scale_xx' and the value is the decoded results
+ optionally with timestamps. `xx` is the ngram LM scale value
+ used during decoding, i.e., 0.1.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ temperature=temperature,
+ )
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+ # at this point, nbest.fsa.scores are all zeros.
+
+ nbest = nbest.intersect(lattice)
+ # Now nbest.fsa.scores contains acoustic scores
+
+ am_scores = nbest.tot_scores()
+
+ # Now we need to compute the LM scores of each path.
+ # (1) Get the token IDs of each Path. We assume the decoding_graph
+ # is an acceptor, i.e., lattice is also an acceptor
+ tokens_shape = nbest.fsa.arcs.shape().remove_axis(1) # [path][arc]
+
+ tokens = k2.RaggedTensor(tokens_shape, nbest.fsa.labels.contiguous())
+ tokens = tokens.remove_values_leq(0) # remove -1 and 0
+
+ token_list: List[List[int]] = tokens.tolist()
+ word_list: List[List[str]] = sp.decode(token_list)
+
+ assert isinstance(oov_word, str), oov_word
+ assert oov_word in word_table, oov_word
+ oov_word_id = word_table[oov_word]
+
+ word_ids_list: List[List[int]] = []
+
+ for words in word_list:
+ this_word_ids = []
+ for w in words.split():
+ if w in word_table:
+ this_word_ids.append(word_table[w])
+ else:
+ this_word_ids.append(oov_word_id)
+ word_ids_list.append(this_word_ids)
+
+ word_fsas = k2.linear_fsa(word_ids_list, device=lattice.device)
+ word_fsas_with_self_loops = k2.add_epsilon_self_loops(word_fsas)
+
+ num_unique_paths = len(word_ids_list)
+
+ b_to_a_map = torch.zeros(
+ num_unique_paths,
+ dtype=torch.int32,
+ device=lattice.device,
+ )
+
+ rescored_word_fsas = k2.intersect_device(
+ a_fsas=G,
+ b_fsas=word_fsas_with_self_loops,
+ b_to_a_map=b_to_a_map,
+ sorted_match_a=True,
+ ret_arc_maps=False,
+ )
+
+ rescored_word_fsas = k2.remove_epsilon_self_loops(rescored_word_fsas)
+ rescored_word_fsas = k2.top_sort(k2.connect(rescored_word_fsas))
+ ngram_lm_scores = rescored_word_fsas.get_tot_scores(
+ use_double_scores=True,
+ log_semiring=False,
+ )
+
+ ans: Dict[str, Union[List[List[int]], DecodingResults]] = {}
+ for s in ngram_lm_scale_list:
+ key = f"ngram_lm_scale_{s}"
+ tot_scores = am_scores.values + s * ngram_lm_scores
+ ragged_tot_scores = k2.RaggedTensor(nbest.shape, tot_scores)
+ max_indexes = ragged_tot_scores.argmax()
+ best_path = k2.index_fsa(nbest.fsa, max_indexes)
+
+ if not return_timestamps:
+ ans[key] = get_texts(best_path)
+ else:
+ ans[key] = get_texts_with_timestamp(best_path)
+
+ return ans
+
+
+def fast_beam_search_with_nbest_rnn_rescoring(
+ model: nn.Module,
+ decoding_graph: k2.Fsa,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ beam: float,
+ max_states: int,
+ max_contexts: int,
+ ngram_lm_scale_list: List[float],
+ num_paths: int,
+ G: k2.Fsa,
+ sp: spm.SentencePieceProcessor,
+ word_table: k2.SymbolTable,
+ rnn_lm_model: torch.nn.Module,
+ rnn_lm_scale_list: List[float],
+ oov_word: str = "",
+ use_double_scores: bool = True,
+ nbest_scale: float = 0.5,
+ temperature: float = 1.0,
+ return_timestamps: bool = False,
+) -> Dict[str, Union[List[List[int]], DecodingResults]]:
+ """It limits the maximum number of symbols per frame to 1.
+ A lattice is first obtained using fast beam search, num_path are selected
+ and rescored using a given language model and a rnn-lm.
+ The shortest path within the lattice is used as the final output.
+
+ Args:
+ model:
+ An instance of `Transducer`.
+ decoding_graph:
+ Decoding graph used for decoding, may be a TrivialGraph or a LG.
+ encoder_out:
+ A tensor of shape (N, T, C) from the encoder.
+ encoder_out_lens:
+ A tensor of shape (N,) containing the number of frames in `encoder_out`
+ before padding.
+ beam:
+ Beam value, similar to the beam used in Kaldi.
+ max_states:
+ Max states per stream per frame.
+ max_contexts:
+ Max contexts pre stream per frame.
+ ngram_lm_scale_list:
+ A list of floats representing LM score scales.
+ num_paths:
+ Number of paths to extract from the decoded lattice.
+ G:
+ An FsaVec containing only a single FSA. It is an n-gram LM.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ rnn_lm_model:
+ A rnn-lm model used for LM rescoring
+ rnn_lm_scale_list:
+ A list of floats representing RNN score scales.
+ oov_word:
+ OOV words are replaced with this word.
+ use_double_scores:
+ True to use double precision for computation. False to use
+ single precision.
+ nbest_scale:
+ It's the scale applied to the lattice.scores. A smaller value
+ yields more unique paths.
+ temperature:
+ Softmax temperature.
+ return_timestamps:
+ Whether to return timestamps.
+ Returns:
+ Return the decoded result in a dict, where the key has the form
+ 'ngram_lm_scale_xx' and the value is the decoded results
+ optionally with timestamps. `xx` is the ngram LM scale value
+ used during decoding, i.e., 0.1.
+ """
+ lattice = fast_beam_search(
+ model=model,
+ decoding_graph=decoding_graph,
+ encoder_out=encoder_out,
+ encoder_out_lens=encoder_out_lens,
+ beam=beam,
+ max_states=max_states,
+ max_contexts=max_contexts,
+ temperature=temperature,
+ )
+
+ nbest = Nbest.from_lattice(
+ lattice=lattice,
+ num_paths=num_paths,
+ use_double_scores=use_double_scores,
+ nbest_scale=nbest_scale,
+ )
+ # at this point, nbest.fsa.scores are all zeros.
+
+ nbest = nbest.intersect(lattice)
+ # Now nbest.fsa.scores contains acoustic scores
+
+ am_scores = nbest.tot_scores()
+
+ # Now we need to compute the LM scores of each path.
+ # (1) Get the token IDs of each Path. We assume the decoding_graph
+ # is an acceptor, i.e., lattice is also an acceptor
+ tokens_shape = nbest.fsa.arcs.shape().remove_axis(1) # [path][arc]
+
+ tokens = k2.RaggedTensor(tokens_shape, nbest.fsa.labels.contiguous())
+ tokens = tokens.remove_values_leq(0) # remove -1 and 0
+
+ token_list: List[List[int]] = tokens.tolist()
+ word_list: List[List[str]] = sp.decode(token_list)
+
+ assert isinstance(oov_word, str), oov_word
+ assert oov_word in word_table, oov_word
+ oov_word_id = word_table[oov_word]
+
+ word_ids_list: List[List[int]] = []
+
+ for words in word_list:
+ this_word_ids = []
+ for w in words.split():
+ if w in word_table:
+ this_word_ids.append(word_table[w])
+ else:
+ this_word_ids.append(oov_word_id)
+ word_ids_list.append(this_word_ids)
+
+ word_fsas = k2.linear_fsa(word_ids_list, device=lattice.device)
+ word_fsas_with_self_loops = k2.add_epsilon_self_loops(word_fsas)
+
+ num_unique_paths = len(word_ids_list)
+
+ b_to_a_map = torch.zeros(
+ num_unique_paths,
+ dtype=torch.int32,
+ device=lattice.device,
+ )
+
+ rescored_word_fsas = k2.intersect_device(
+ a_fsas=G,
+ b_fsas=word_fsas_with_self_loops,
+ b_to_a_map=b_to_a_map,
+ sorted_match_a=True,
+ ret_arc_maps=False,
+ )
+
+ rescored_word_fsas = k2.remove_epsilon_self_loops(rescored_word_fsas)
+ rescored_word_fsas = k2.top_sort(k2.connect(rescored_word_fsas))
+ ngram_lm_scores = rescored_word_fsas.get_tot_scores(
+ use_double_scores=True,
+ log_semiring=False,
+ )
+
+ # Now RNN-LM
+ blank_id = model.decoder.blank_id
+ sos_id = sp.piece_to_id("sos_id")
+ eos_id = sp.piece_to_id("eos_id")
+
+ sos_tokens = add_sos(tokens, sos_id)
+ tokens_eos = add_eos(tokens, eos_id)
+ sos_tokens_row_splits = sos_tokens.shape.row_splits(1)
+ sentence_lengths = sos_tokens_row_splits[1:] - sos_tokens_row_splits[:-1]
+
+ x_tokens = sos_tokens.pad(mode="constant", padding_value=blank_id)
+ y_tokens = tokens_eos.pad(mode="constant", padding_value=blank_id)
+
+ x_tokens = x_tokens.to(torch.int64)
+ y_tokens = y_tokens.to(torch.int64)
+ sentence_lengths = sentence_lengths.to(torch.int64)
+
+ rnn_lm_nll = rnn_lm_model(x=x_tokens, y=y_tokens, lengths=sentence_lengths)
+ assert rnn_lm_nll.ndim == 2
+ assert rnn_lm_nll.shape[0] == len(token_list)
+ rnn_lm_scores = -1 * rnn_lm_nll.sum(dim=1)
+
+ ans: Dict[str, List[List[int]]] = {}
+ for n_scale in ngram_lm_scale_list:
+ for rnn_scale in rnn_lm_scale_list:
+ key = f"ngram_lm_scale_{n_scale}_rnn_lm_scale_{rnn_scale}"
+ tot_scores = (
+ am_scores.values + n_scale * ngram_lm_scores + rnn_scale * rnn_lm_scores
+ )
+ ragged_tot_scores = k2.RaggedTensor(nbest.shape, tot_scores)
+ max_indexes = ragged_tot_scores.argmax()
+ best_path = k2.index_fsa(nbest.fsa, max_indexes)
+
+ if not return_timestamps:
+ ans[key] = get_texts(best_path)
+ else:
+ ans[key] = get_texts_with_timestamp(best_path)
+
+ return ans
+
+
+def modified_beam_search_ngram_rescoring(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ ngram_lm: NgramLm,
+ ngram_lm_scale: float,
+ beam: int = 4,
+ temperature: float = 1.0,
+) -> List[List[int]]:
+ """Beam search in batch mode with --max-sym-per-frame=1 being hardcoded.
+
+ Args:
+ model:
+ The transducer model.
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, C).
+ encoder_out_lens:
+ A 1-D tensor of shape (N,), containing number of valid frames in
+ encoder_out before padding.
+ beam:
+ Number of active paths during the beam search.
+ temperature:
+ Softmax temperature.
+ Returns:
+ Return a list-of-list of token IDs. ans[i] is the decoding results
+ for the i-th utterance.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+ lm_scale = ngram_lm_scale
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ state_cost=NgramLmStateCost(ngram_lm),
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for batch_size in batch_size_list:
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end]
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [
+ hyp.log_prob.reshape(1, 1) + hyp.state_cost.lm_score * lm_scale
+ for hyps in A
+ for hyp in hyps
+ ]
+ ) # (num_hyps, 1)
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+ # decoder_out is of shape (num_hyps, 1, 1, joiner_dim)
+
+ # Note: For torch 1.7.1 and below, it requires a torch.int64 tensor
+ # as index, so we use `to(torch.int64)` below.
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = (logits / temperature).log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+ vocab_size = log_probs.size(-1)
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_ys = hyp.ys[:]
+ new_token = topk_token_indexes[k]
+ if new_token not in (blank_id, unk_id):
+ new_ys.append(new_token)
+ state_cost = hyp.state_cost.forward_one_step(new_token)
+ else:
+ state_cost = hyp.state_cost
+
+ # We only keep AM scores in new_hyp.log_prob
+ new_log_prob = topk_log_probs[k] - hyp.state_cost.lm_score * lm_scale
+
+ new_hyp = Hypothesis(
+ ys=new_ys, log_prob=new_log_prob, state_cost=state_cost
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+ best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+ sorted_ans = [h.ys[context_size:] for h in best_hyps]
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+
+ return ans
+
+
+def modified_beam_search_LODR(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ LODR_lm: NgramLm,
+ LODR_lm_scale: float,
+ LM: LmScorer,
+ beam: int = 4,
+ context_graph: Optional[ContextGraph] = None,
+) -> List[List[int]]:
+ """This function implements LODR (https://arxiv.org/abs/2203.16776) with
+ `modified_beam_search`. It uses a bi-gram language model as the estimate
+ of the internal language model and subtracts its score during shallow fusion
+ with an external language model. This implementation uses a RNNLM as the
+ external language model.
+
+ Args:
+ model (Transducer):
+ The transducer model
+ encoder_out (torch.Tensor):
+ Encoder output in (N,T,C)
+ encoder_out_lens (torch.Tensor):
+ A 1-D tensor of shape (N,), containing the number of
+ valid frames in encoder_out before padding.
+ LODR_lm:
+ A low order n-gram LM, whose score will be subtracted during shallow fusion
+ LODR_lm_scale:
+ The scale of the LODR_lm
+ LM:
+ A neural net LM, e.g an RNNLM or transformer LM
+ beam (int, optional):
+ Beam size. Defaults to 4.
+
+ Returns:
+ Return a list-of-list of token IDs. ans[i] is the decoding results
+ for the i-th utterance.
+
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+ assert LM is not None
+ lm_scale = LM.lm_scale
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ sos_id = getattr(LM, "sos_id", 1)
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ # get initial lm score and lm state by scoring the "sos" token
+ sos_token = torch.tensor([[sos_id]]).to(torch.int64).to(device)
+ lens = torch.tensor([1]).to(device)
+ init_score, init_states = LM.score_token(sos_token, lens)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ state=init_states, # state of the NN LM
+ lm_score=init_score.reshape(-1),
+ state_cost=NgramLmStateCost(
+ LODR_lm
+ ), # state of the source domain ngram
+ context_state=None if context_graph is None else context_graph.root,
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for batch_size in batch_size_list:
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end] # get batch
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ )
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = logits.log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+ """
+ for all hyps with a non-blank new token, score this token.
+ It is a little confusing here because this for-loop
+ looks very similar to the one below. Here, we go through all
+ top-k tokens and only add the non-blanks ones to the token_list.
+ LM will score those tokens given the LM states. Note that
+ the variable `scores` is the LM score after seeing the new
+ non-blank token.
+ """
+ token_list = []
+ hs = []
+ cs = []
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_token = topk_token_indexes[k]
+ if new_token not in (blank_id, unk_id):
+ if LM.lm_type == "rnn":
+ token_list.append([new_token])
+ # store the LSTM states
+ hs.append(hyp.state[0])
+ cs.append(hyp.state[1])
+ else:
+ # for transformer LM
+ token_list.append(
+ [sos_id] + hyp.ys[context_size:] + [new_token]
+ )
+
+ # forward NN LM to get new states and scores
+ if len(token_list) != 0:
+ x_lens = torch.tensor([len(tokens) for tokens in token_list]).to(device)
+ if LM.lm_type == "rnn":
+ tokens_to_score = (
+ torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
+ )
+ hs = torch.cat(hs, dim=1).to(device)
+ cs = torch.cat(cs, dim=1).to(device)
+ state = (hs, cs)
+ else:
+ # for transformer LM
+ tokens_list = [torch.tensor(tokens) for tokens in token_list]
+ tokens_to_score = (
+ torch.nn.utils.rnn.pad_sequence(
+ tokens_list, batch_first=True, padding_value=0.0
+ )
+ .to(device)
+ .to(torch.int64)
+ )
+
+ state = None
+
+ scores, lm_states = LM.score_token(tokens_to_score, x_lens, state)
+
+ count = 0 # index, used to locate score and lm states
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ ys = hyp.ys[:]
+
+ # current score of hyp
+ lm_score = hyp.lm_score
+ state = hyp.state
+
+ hyp_log_prob = topk_log_probs[k] # get score of current hyp
+ new_token = topk_token_indexes[k]
+
+ context_score = 0
+ new_context_state = None if context_graph is None else hyp.context_state
+ if new_token not in (blank_id, unk_id):
+ if context_graph is not None:
+ (
+ context_score,
+ new_context_state,
+ ) = context_graph.forward_one_step(hyp.context_state, new_token)
+
+ ys.append(new_token)
+ state_cost = hyp.state_cost.forward_one_step(new_token)
+
+ # calculate the score of the latest token
+ current_ngram_score = state_cost.lm_score - hyp.state_cost.lm_score
+
+ assert current_ngram_score <= 0.0, (
+ state_cost.lm_score,
+ hyp.state_cost.lm_score,
+ )
+ # score = score + TDLM_score - LODR_score
+ # LODR_LM_scale should be a negative number here
+ hyp_log_prob += (
+ lm_score[new_token] * lm_scale
+ + LODR_lm_scale * current_ngram_score
+ + context_score
+ ) # add the lm score
+
+ lm_score = scores[count]
+ if LM.lm_type == "rnn":
+ state = (
+ lm_states[0][:, count, :].unsqueeze(1),
+ lm_states[1][:, count, :].unsqueeze(1),
+ )
+ count += 1
+ else:
+ state_cost = hyp.state_cost
+
+ new_hyp = Hypothesis(
+ ys=ys,
+ log_prob=hyp_log_prob,
+ state=state,
+ lm_score=lm_score,
+ state_cost=state_cost,
+ context_state=new_context_state,
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+
+ # finalize context_state, if the matched contexts do not reach final state
+ # we need to add the score on the corresponding backoff arc
+ if context_graph is not None:
+ finalized_B = [HypothesisList() for _ in range(len(B))]
+ for i, hyps in enumerate(B):
+ for hyp in list(hyps):
+ context_score, new_context_state = context_graph.finalize(
+ hyp.context_state
+ )
+ finalized_B[i].add(
+ Hypothesis(
+ ys=hyp.ys,
+ log_prob=hyp.log_prob + context_score,
+ timestamp=hyp.timestamp,
+ context_state=new_context_state,
+ )
+ )
+ B = finalized_B
+
+ best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+ sorted_ans = [h.ys[context_size:] for h in best_hyps]
+ ans = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+
+ return ans
+
+
+def modified_beam_search_lm_shallow_fusion(
+ model: nn.Module,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ LM: LmScorer,
+ beam: int = 4,
+ return_timestamps: bool = False,
+) -> List[List[int]]:
+ """Modified_beam_search + NN LM shallow fusion
+
+ Args:
+ model (Transducer):
+ The transducer model
+ encoder_out (torch.Tensor):
+ Encoder output in (N,T,C)
+ encoder_out_lens (torch.Tensor):
+ A 1-D tensor of shape (N,), containing the number of
+ valid frames in encoder_out before padding.
+ sp:
+ Sentence piece generator.
+ LM (LmScorer):
+ A neural net LM, e.g RNN or Transformer
+ beam (int, optional):
+ Beam size. Defaults to 4.
+
+ Returns:
+ Return a list-of-list of token IDs. ans[i] is the decoding results
+ for the i-th utterance.
+ """
+ assert encoder_out.ndim == 3, encoder_out.shape
+ assert encoder_out.size(0) >= 1, encoder_out.size(0)
+ assert LM is not None
+ lm_scale = LM.lm_scale
+
+ packed_encoder_out = torch.nn.utils.rnn.pack_padded_sequence(
+ input=encoder_out,
+ lengths=encoder_out_lens.cpu(),
+ batch_first=True,
+ enforce_sorted=False,
+ )
+
+ blank_id = model.decoder.blank_id
+ sos_id = getattr(LM, "sos_id", 1)
+ unk_id = getattr(model, "unk_id", blank_id)
+ context_size = model.decoder.context_size
+ device = next(model.parameters()).device
+
+ batch_size_list = packed_encoder_out.batch_sizes.tolist()
+ N = encoder_out.size(0)
+ assert torch.all(encoder_out_lens > 0), encoder_out_lens
+ assert N == batch_size_list[0], (N, batch_size_list)
+
+ # get initial lm score and lm state by scoring the "sos" token
+ sos_token = torch.tensor([[sos_id]]).to(torch.int64).to(device)
+ lens = torch.tensor([1]).to(device)
+ init_score, init_states = LM.score_token(sos_token, lens)
+
+ B = [HypothesisList() for _ in range(N)]
+ for i in range(N):
+ B[i].add(
+ Hypothesis(
+ ys=[-1] * (context_size - 1) + [blank_id],
+ log_prob=torch.zeros(1, dtype=torch.float32, device=device),
+ state=init_states,
+ lm_score=init_score.reshape(-1),
+ timestamp=[],
+ )
+ )
+
+ encoder_out = model.joiner.encoder_proj(packed_encoder_out.data)
+
+ offset = 0
+ finalized_B = []
+ for t, batch_size in enumerate(batch_size_list):
+ start = offset
+ end = offset + batch_size
+ current_encoder_out = encoder_out.data[start:end] # get batch
+ current_encoder_out = current_encoder_out.unsqueeze(1).unsqueeze(1)
+ # current_encoder_out's shape is (batch_size, 1, 1, encoder_out_dim)
+ offset = end
+
+ finalized_B = B[batch_size:] + finalized_B
+ B = B[:batch_size]
+
+ hyps_shape = get_hyps_shape(B).to(device)
+
+ A = [list(b) for b in B]
+ B = [HypothesisList() for _ in range(batch_size)]
+
+ ys_log_probs = torch.cat(
+ [hyp.log_prob.reshape(1, 1) for hyps in A for hyp in hyps]
+ )
+
+ lm_scores = torch.cat(
+ [hyp.lm_score.reshape(1, -1) for hyps in A for hyp in hyps]
+ )
+
+ decoder_input = torch.tensor(
+ [hyp.ys[-context_size:] for hyps in A for hyp in hyps],
+ device=device,
+ dtype=torch.int64,
+ ) # (num_hyps, context_size)
+
+ decoder_out = model.decoder(decoder_input, need_pad=False).unsqueeze(1)
+ decoder_out = model.joiner.decoder_proj(decoder_out)
+
+ current_encoder_out = torch.index_select(
+ current_encoder_out,
+ dim=0,
+ index=hyps_shape.row_ids(1).to(torch.int64),
+ ) # (num_hyps, 1, 1, encoder_out_dim)
+
+ logits = model.joiner(
+ current_encoder_out,
+ decoder_out,
+ project_input=False,
+ ) # (num_hyps, 1, 1, vocab_size)
+
+ logits = logits.squeeze(1).squeeze(1) # (num_hyps, vocab_size)
+
+ log_probs = logits.log_softmax(dim=-1) # (num_hyps, vocab_size)
+
+ log_probs.add_(ys_log_probs)
+
+ vocab_size = log_probs.size(-1)
+
+ log_probs = log_probs.reshape(-1)
+
+ row_splits = hyps_shape.row_splits(1) * vocab_size
+ log_probs_shape = k2.ragged.create_ragged_shape2(
+ row_splits=row_splits, cached_tot_size=log_probs.numel()
+ )
+ ragged_log_probs = k2.RaggedTensor(shape=log_probs_shape, value=log_probs)
+ """
+ for all hyps with a non-blank new token, score this token.
+ It is a little confusing here because this for-loop
+ looks very similar to the one below. Here, we go through all
+ top-k tokens and only add the non-blanks ones to the token_list.
+ `LM` will score those tokens given the LM states. Note that
+ the variable `scores` is the LM score after seeing the new
+ non-blank token.
+ """
+ token_list = [] # a list of list
+ hs = []
+ cs = []
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ new_token = topk_token_indexes[k]
+ if new_token not in (blank_id, unk_id):
+ if LM.lm_type == "rnn":
+ token_list.append([new_token])
+ # store the LSTM states
+ hs.append(hyp.state[0])
+ cs.append(hyp.state[1])
+ else:
+ # for transformer LM
+ token_list.append(
+ [sos_id] + hyp.ys[context_size:] + [new_token]
+ )
+
+ if len(token_list) != 0:
+ x_lens = torch.tensor([len(tokens) for tokens in token_list]).to(device)
+ if LM.lm_type == "rnn":
+ tokens_to_score = (
+ torch.tensor(token_list).to(torch.int64).to(device).reshape(-1, 1)
+ )
+ hs = torch.cat(hs, dim=1).to(device)
+ cs = torch.cat(cs, dim=1).to(device)
+ state = (hs, cs)
+ else:
+ # for transformer LM
+ tokens_list = [torch.tensor(tokens) for tokens in token_list]
+ tokens_to_score = (
+ torch.nn.utils.rnn.pad_sequence(
+ tokens_list, batch_first=True, padding_value=0.0
+ )
+ .to(device)
+ .to(torch.int64)
+ )
+
+ state = None
+
+ scores, lm_states = LM.score_token(tokens_to_score, x_lens, state)
+
+ count = 0 # index, used to locate score and lm states
+ for i in range(batch_size):
+ topk_log_probs, topk_indexes = ragged_log_probs[i].topk(beam)
+
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ topk_hyp_indexes = (topk_indexes // vocab_size).tolist()
+ topk_token_indexes = (topk_indexes % vocab_size).tolist()
+
+ for k in range(len(topk_hyp_indexes)):
+ hyp_idx = topk_hyp_indexes[k]
+ hyp = A[i][hyp_idx]
+
+ ys = hyp.ys[:]
+
+ lm_score = hyp.lm_score
+ state = hyp.state
+
+ hyp_log_prob = topk_log_probs[k] # get score of current hyp
+ new_token = topk_token_indexes[k]
+ new_timestamp = hyp.timestamp[:]
+ if new_token not in (blank_id, unk_id):
+ ys.append(new_token)
+ new_timestamp.append(t)
+
+ hyp_log_prob += lm_score[new_token] * lm_scale # add the lm score
+
+ lm_score = scores[count]
+ if LM.lm_type == "rnn":
+ state = (
+ lm_states[0][:, count, :].unsqueeze(1),
+ lm_states[1][:, count, :].unsqueeze(1),
+ )
+ count += 1
+
+ new_hyp = Hypothesis(
+ ys=ys,
+ log_prob=hyp_log_prob,
+ state=state,
+ lm_score=lm_score,
+ timestamp=new_timestamp,
+ )
+ B[i].add(new_hyp)
+
+ B = B + finalized_B
+ best_hyps = [b.get_most_probable(length_norm=True) for b in B]
+
+ sorted_ans = [h.ys[context_size:] for h in best_hyps]
+ sorted_timestamps = [h.timestamp for h in best_hyps]
+ ans = []
+ ans_timestamps = []
+ unsorted_indices = packed_encoder_out.unsorted_indices.tolist()
+ for i in range(N):
+ ans.append(sorted_ans[unsorted_indices[i]])
+ ans_timestamps.append(sorted_timestamps[unsorted_indices[i]])
+
+ if not return_timestamps:
+ return ans
+ else:
+ return DecodingResults(
+ hyps=ans,
+ timestamps=ans_timestamps,
+ )
diff --git a/egs/europarl_st/SRT/lcma_srt/datamodule.py b/egs/europarl_st/SRT/lcma_srt/datamodule.py
new file mode 100644
index 0000000000..ac15f223ac
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/datamodule.py
@@ -0,0 +1,583 @@
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import inspect
+import logging
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import torch
+import torch.distributed as dist
+from lhotse import CutSet, Fbank, FbankConfig, load_manifest, load_manifest_lazy
+from lhotse.dataset import ( # noqa F401 for PrecomputedFeatures
+ CutConcatenate,
+ CutMix,
+ DynamicBucketingSampler,
+ K2SpeechRecognitionDataset,
+ PrecomputedFeatures,
+ SimpleCutSampler,
+ SpecAugment,
+)
+from lhotse.dataset.input_strategies import ( # noqa F401 For AudioSamples
+ AudioSamples,
+ OnTheFlyFeatures,
+)
+from lhotse.utils import fix_random_seed
+from torch.distributed import get_rank, get_world_size, is_initialized
+from torch.utils.data import DataLoader
+
+from icefall.utils import str2bool
+
+
+class _SeedWorkers:
+ def __init__(self, seed: int):
+ self.seed = seed
+
+ def __call__(self, worker_id: int):
+ fix_random_seed(self.seed + worker_id)
+
+
+class LibriSpeechAsrDataModule:
+ """
+ DataModule for k2 ASR experiments.
+ It assumes there is always one train and valid dataloader,
+ but there can be multiple test dataloaders (e.g. LibriSpeech test-clean
+ and test-other).
+
+ It contains all the common data pipeline modules used in ASR
+ experiments, e.g.:
+ - dynamic batch size,
+ - bucketing samplers,
+ - cut concatenation,
+ - augmentation,
+ - on-the-fly feature extraction
+
+ This class should be derived for specific corpora used in ASR tasks.
+ """
+
+ def __init__(self, args: argparse.Namespace):
+ self.args = args
+
+ @classmethod
+ def add_arguments(cls, parser: argparse.ArgumentParser):
+ group = parser.add_argument_group(
+ title="ASR data related options",
+ description="These options are used for the preparation of "
+ "PyTorch DataLoaders from Lhotse CutSet's -- they control the "
+ "effective batch sizes, sampling strategies, applied data "
+ "augmentations, etc.",
+ )
+ group.add_argument(
+ "--full-libri",
+ type=str2bool,
+ default=True,
+ help="""Used only when --mini-libri is False.When enabled,
+ use 960h LibriSpeech. Otherwise, use 100h subset.""",
+ )
+ group.add_argument(
+ "--mini-libri",
+ type=str2bool,
+ default=False,
+ help="True for mini librispeech",
+ )
+
+ group.add_argument(
+ "--manifest-dir",
+ type=Path,
+ default=Path("fbank"),
+ help="Path to directory with train/valid/test cuts.",
+ )
+ group.add_argument(
+ "--max-duration",
+ type=int,
+ default=200.0,
+ help="Maximum pooled recordings duration (seconds) in a "
+ "single batch. You can reduce it if it causes CUDA OOM.",
+ )
+ group.add_argument(
+ "--bucketing-sampler",
+ type=str2bool,
+ default=True,
+ help="When enabled, the batches will come from buckets of "
+ "similar duration (saves padding frames).",
+ )
+ group.add_argument(
+ "--num-buckets",
+ type=int,
+ default=30,
+ help="The number of buckets for the DynamicBucketingSampler"
+ "(you might want to increase it for larger datasets).",
+ )
+ group.add_argument(
+ "--concatenate-cuts",
+ type=str2bool,
+ default=False,
+ help="When enabled, utterances (cuts) will be concatenated "
+ "to minimize the amount of padding.",
+ )
+ group.add_argument(
+ "--duration-factor",
+ type=float,
+ default=1.0,
+ help="Determines the maximum duration of a concatenated cut "
+ "relative to the duration of the longest cut in a batch.",
+ )
+ group.add_argument(
+ "--gap",
+ type=float,
+ default=1.0,
+ help="The amount of padding (in seconds) inserted between "
+ "concatenated cuts. This padding is filled with noise when "
+ "noise augmentation is used.",
+ )
+ group.add_argument(
+ "--on-the-fly-feats",
+ type=str2bool,
+ default=False,
+ help="When enabled, use on-the-fly cut mixing and feature "
+ "extraction. Will drop existing precomputed feature manifests "
+ "if available.",
+ )
+ group.add_argument(
+ "--shuffle",
+ type=str2bool,
+ default=True,
+ help="When enabled (=default), the examples will be "
+ "shuffled for each epoch.",
+ )
+ group.add_argument(
+ "--drop-last",
+ type=str2bool,
+ default=True,
+ help="Whether to drop last batch. Used by sampler.",
+ )
+ group.add_argument(
+ "--return-cuts",
+ type=str2bool,
+ default=True,
+ help="When enabled, each batch will have the "
+ "field: batch['supervisions']['cut'] with the cuts that "
+ "were used to construct it.",
+ )
+
+ group.add_argument(
+ "--num-workers",
+ type=int,
+ default=2,
+ help="The number of training dataloader workers that "
+ "collect the batches.",
+ )
+
+ group.add_argument(
+ "--enable-spec-aug",
+ type=str2bool,
+ default=True,
+ help="When enabled, use SpecAugment for training dataset.",
+ )
+
+ group.add_argument(
+ "--spec-aug-time-warp-factor",
+ type=int,
+ default=80,
+ help="Used only when --enable-spec-aug is True. "
+ "It specifies the factor for time warping in SpecAugment. "
+ "Larger values mean more warping. "
+ "A value less than 1 means to disable time warp.",
+ )
+
+ group.add_argument(
+ "--enable-musan",
+ type=str2bool,
+ default=True,
+ help="When enabled, select noise from MUSAN and mix it"
+ "with training dataset. ",
+ )
+
+ group.add_argument(
+ "--input-strategy",
+ type=str,
+ default="PrecomputedFeatures",
+ help="AudioSamples or PrecomputedFeatures",
+ )
+
+ group.add_argument(
+ "--dev-name",
+ type=str,
+ help="Path to the validation/dev dataset manifest file.",
+ )
+ group.add_argument(
+ "--test-name",
+ type=str,
+ help="Path to the test dataset manifest file.",
+ )
+
+ def train_dataloaders(
+ self,
+ cuts_train: CutSet,
+ sampler_state_dict: Optional[Dict[str, Any]] = None,
+ ) -> DataLoader:
+ """
+ Args:
+ cuts_train:
+ CutSet for training.
+ sampler_state_dict:
+ The state dict for the training sampler.
+ """
+ # 1. Initialize distributed parameters
+ rank = get_rank() if is_initialized() else 0
+ world_size = get_world_size() if is_initialized() else 1
+ is_distributed = is_initialized()
+
+ # 2. Print key configuration info
+ logging.info(
+ f"\n{'='*50}\n"
+ f"Initializing dataloader:\n"
+ f" - Distributed mode: {'ON' if is_distributed else 'OFF'}\n"
+ f" - Rank/World size: {rank}/{world_size}\n"
+ f" - Max duration per GPU: {self.args.max_duration}s\n"
+ f"{'='*50}"
+ )
+ transforms = []
+ if self.args.enable_musan:
+ logging.info("Enable MUSAN")
+ logging.info("About to get Musan cuts")
+ cuts_musan = load_manifest(self.args.manifest_dir / "musan_cuts.jsonl.gz")
+ transforms.append(
+ CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
+ )
+ else:
+ logging.info("Disable MUSAN")
+
+ if self.args.concatenate_cuts:
+ logging.info(
+ f"Using cut concatenation with duration factor "
+ f"{self.args.duration_factor} and gap {self.args.gap}."
+ )
+ # Cut concatenation should be the first transform in the list,
+ # so that if we e.g. mix noise in, it will fill the gaps between
+ # different utterances.
+ transforms = [
+ CutConcatenate(
+ duration_factor=self.args.duration_factor, gap=self.args.gap
+ )
+ ] + transforms
+
+ input_transforms = []
+ if self.args.enable_spec_aug:
+ logging.info("Enable SpecAugment")
+ logging.info(f"Time warp factor: {self.args.spec_aug_time_warp_factor}")
+ # Set the value of num_frame_masks according to Lhotse's version.
+ # In different Lhotse's versions, the default of num_frame_masks is
+ # different.
+ num_frame_masks = 10
+ num_frame_masks_parameter = inspect.signature(
+ SpecAugment.__init__
+ ).parameters["num_frame_masks"]
+ if num_frame_masks_parameter.default == 1:
+ num_frame_masks = 2
+ logging.info(f"Num frame mask: {num_frame_masks}")
+ input_transforms.append(
+ SpecAugment(
+ time_warp_factor=self.args.spec_aug_time_warp_factor,
+ num_frame_masks=num_frame_masks,
+ features_mask_size=27,
+ num_feature_masks=2,
+ frames_mask_size=100,
+ )
+ )
+ else:
+ logging.info("Disable SpecAugment")
+
+ logging.info("About to create train dataset")
+ train = K2SpeechRecognitionDataset(
+ input_strategy=eval(self.args.input_strategy)(),
+ cut_transforms=transforms,
+ input_transforms=input_transforms,
+ return_cuts=self.args.return_cuts,
+ )
+
+ if self.args.on_the_fly_feats:
+ # NOTE: the PerturbSpeed transform should be added only if we
+ # remove it from data prep stage.
+ # Add on-the-fly speed perturbation; since originally it would
+ # have increased epoch size by 3, we will apply prob 2/3 and use
+ # 3x more epochs.
+ # Speed perturbation probably should come first before
+ # concatenation, but in principle the transforms order doesn't have
+ # to be strict (e.g. could be randomized)
+ # transforms = [PerturbSpeed(factors=[0.9, 1.1], p=2/3)] + transforms # noqa
+ # Drop feats to be on the safe side.
+ train = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
+ input_transforms=input_transforms,
+ return_cuts=self.args.return_cuts,
+ )
+
+ if self.args.bucketing_sampler:
+ logging.info("Using DynamicBucketingSampler.")
+ train_sampler = DynamicBucketingSampler(
+ cuts_train,
+ max_duration=self.args.max_duration,
+ shuffle=self.args.shuffle,
+ num_buckets=self.args.num_buckets,
+ buffer_size=self.args.num_buckets * 2000,
+ shuffle_buffer_size=self.args.num_buckets * 5000,
+ rank=rank,
+ world_size=world_size,
+ drop_last=self.args.drop_last,
+ )
+ sampler_type = "DynamicBucketingSampler"
+
+ else:
+ logging.info("Using SimpleCutSampler.")
+ train_sampler = SimpleCutSampler(
+ cuts_train,
+ max_duration=self.args.max_duration,
+ shuffle=self.args.shuffle,
+ )
+ logging.info("About to create train dataloader")
+
+ if sampler_state_dict is not None:
+ logging.info("Loading sampler state dict")
+ train_sampler.load_state_dict(sampler_state_dict)
+
+ # 'seed' is derived from the current random state, which will have
+ # previously been set in the main process.
+ seed = torch.randint(0, 100000, ()).item()
+ worker_init_fn = _SeedWorkers(seed)
+
+ train_dl = DataLoader(
+ train,
+ sampler=train_sampler,
+ batch_size=None,
+ num_workers=self.args.num_workers,
+ persistent_workers=False,
+ worker_init_fn=worker_init_fn,
+ )
+
+ return train_dl
+
+ def valid_dataloaders_gpt41(self, cuts_valid: CutSet) -> DataLoader[Any]:
+ """
+ Create validation DataLoader with optional distributed training support.
+
+ Args:
+ cuts_valid: CutSet containing validation data.
+
+ Returns:
+ Configured DataLoader for validation.
+ """
+ transforms = []
+ if self.args.concatenate_cuts:
+ transforms = [
+ CutConcatenate(
+ duration_factor=self.args.duration_factor, gap=self.args.gap
+ )
+ ] + transforms
+
+ logging.info("## Creating validation dataset")
+ if self.args.on_the_fly_feats:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
+ return_cuts=self.args.return_cuts,
+ )
+ else:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ return_cuts=self.args.return_cuts,
+ )
+
+ world_size = dist.get_world_size() if dist.is_initialized() else 1
+ rank = dist.get_rank() if dist.is_initialized() else 0
+
+ logging.info(f"## Creating sampler (world_size={world_size}, rank={rank})")
+ valid_sampler = DynamicBucketingSampler(
+ cuts_valid,
+ max_duration=self.args.max_duration,
+ shuffle=False,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ logging.info("## Creating validation DataLoader")
+ valid_dl = DataLoader(
+ validate,
+ sampler=valid_sampler,
+ batch_size=None,
+ num_workers=self.args.num_workers
+ if hasattr(self.args, "num_workers")
+ else 2,
+ persistent_workers=False,
+ )
+
+ return valid_dl
+
+ def valid_dataloaders(self, cuts_valid: CutSet) -> DataLoader:
+ transforms = []
+ if self.args.concatenate_cuts:
+ transforms = [
+ CutConcatenate(
+ duration_factor=self.args.duration_factor, gap=self.args.gap
+ )
+ ] + transforms
+
+ logging.info("About to create dev dataset")
+ if self.args.on_the_fly_feats:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
+ return_cuts=self.args.return_cuts,
+ )
+ else:
+ validate = K2SpeechRecognitionDataset(
+ cut_transforms=transforms,
+ return_cuts=self.args.return_cuts,
+ )
+ valid_sampler = DynamicBucketingSampler(
+ cuts_valid,
+ max_duration=self.args.max_duration,
+ shuffle=False,
+ )
+ logging.info("About to create dev dataloader")
+ valid_dl = DataLoader(
+ validate,
+ sampler=valid_sampler,
+ batch_size=None,
+ num_workers=2,
+ persistent_workers=False,
+ )
+
+ return valid_dl
+
+ def test_dataloaders(self, cuts: CutSet) -> DataLoader:
+ logging.debug("About to create test dataset")
+ test = K2SpeechRecognitionDataset(
+ input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80)))
+ if self.args.on_the_fly_feats
+ else eval(self.args.input_strategy)(),
+ return_cuts=self.args.return_cuts,
+ )
+ sampler = DynamicBucketingSampler(
+ cuts,
+ max_duration=self.args.max_duration,
+ shuffle=False,
+ )
+ logging.debug("About to create test dataloader")
+ test_dl = DataLoader(
+ test,
+ batch_size=None,
+ sampler=sampler,
+ num_workers=self.args.num_workers,
+ )
+ return test_dl
+
+ @lru_cache()
+ def train_clean_5_cuts(self) -> CutSet:
+ logging.info("mini_librispeech: About to get train-clean-5 cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_train-clean-5.jsonl.gz"
+ )
+
+ @lru_cache()
+ def train_clean_100_cuts(self) -> CutSet:
+ logging.info("About to get train-clean-100 cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_train-clean-100.jsonl.gz"
+ )
+
+ @lru_cache()
+ def train_clean_360_cuts(self) -> CutSet:
+ logging.info("About to get train-clean-360 cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_train-clean-360.jsonl.gz"
+ )
+
+ @lru_cache()
+ def train_other_500_cuts(self) -> CutSet:
+ logging.info("About to get train-other-500 cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_train-other-500.jsonl.gz"
+ )
+
+ @lru_cache()
+ def train_all_shuf_cuts(self) -> CutSet:
+ logging.info(
+ "About to get the shuffled train-clean-100, \
+ train-clean-360 and train-other-500 cuts"
+ )
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_train-all-shuf.jsonl.gz"
+ )
+
+ @lru_cache()
+ def dev_clean_2_cuts(self) -> CutSet:
+ logging.info("mini_librispeech: About to get dev-clean-2 cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_dev-clean-2.jsonl.gz"
+ )
+
+ @lru_cache()
+ def dev_clean_cuts(self) -> CutSet:
+ logging.info("About to get dev cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_dev-clean.jsonl.gz"
+ )
+
+ @lru_cache()
+ def dev_other_cuts(self) -> CutSet:
+ logging.info("About to get dev-other cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_dev-other.jsonl.gz"
+ )
+
+ @lru_cache()
+ def test_clean_cuts(self) -> CutSet:
+ logging.info("About to get test cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_test-clean.jsonl.gz"
+ )
+
+ @lru_cache()
+ def test_other_cuts(self) -> CutSet:
+ logging.info("About to get test-other cuts")
+ return load_manifest_lazy(
+ self.args.manifest_dir / "librispeech_cuts_test-other.jsonl.gz"
+ )
+
+ @lru_cache()
+ def gigaspeech_subset_small_cuts(self) -> CutSet:
+ logging.info("About to get Gigaspeech subset-S cuts")
+ return load_manifest_lazy(self.args.manifest_dir / "cuts_S.jsonl.gz")
+
+ @lru_cache()
+ def gigaspeech_dev_cuts(self) -> CutSet:
+ logging.info("About to get Gigaspeech dev cuts")
+ return load_manifest_lazy(self.args.manifest_dir / "cuts_DEV.jsonl.gz")
+
+ @lru_cache()
+ def dev_cuts(self) -> CutSet:
+ logging.info("About to get dev cuts")
+ return load_manifest_lazy(self.args.manifest_dir / self.args.dev_name)
+
+ @lru_cache()
+ def test_cuts(self) -> CutSet:
+ logging.info("About to get test cuts")
+ test_path = Path(self.args.test_name)
+ if test_path.exists():
+ return load_manifest_lazy(test_path)
+ return load_manifest_lazy(self.args.manifest_dir / self.args.test_name)
diff --git a/egs/europarl_st/SRT/lcma_srt/decode.py b/egs/europarl_st/SRT/lcma_srt/decode.py
new file mode 100644
index 0000000000..38cf829035
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode.py
@@ -0,0 +1,1391 @@
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import logging
+import math
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
+import k2
+import sentencepiece as spm
+import torch
+import torch.nn as nn
+from beam_search import (
+ beam_search,
+ fast_beam_search_nbest,
+ fast_beam_search_nbest_LG,
+ fast_beam_search_nbest_oracle,
+ fast_beam_search_one_best,
+ greedy_search,
+ greedy_search_batch,
+ modified_beam_search,
+ modified_beam_search_lm_rescore,
+ modified_beam_search_lm_rescore_LODR,
+ modified_beam_search_lm_shallow_fusion,
+ modified_beam_search_LODR,
+)
+from datamodule import LibriSpeechAsrDataModule
+from lhotse import set_caching_enabled
+from lhotse.cut import Cut
+from train_cross_node_jsrt import (
+ add_model_arguments,
+ build_srctgt_lang_list,
+ get_model,
+ get_params,
+)
+
+from icefall import ContextGraph, LmScorer, NgramLm
+from icefall.checkpoint import (
+ average_checkpoints,
+ average_checkpoints_with_averaged_model,
+ find_checkpoints,
+ load_checkpoint,
+)
+from icefall.lexicon import Lexicon
+from icefall.utils import (
+ AttributeDict,
+ setup_logger,
+ store_transcripts,
+ str2bool,
+ write_error_stats,
+)
+
+LOG_EPS = math.log(1e-10)
+
+
+def _normalize_lang_tag(tag: Optional[str]) -> Optional[str]:
+ if tag is None:
+ return None
+ if not isinstance(tag, str):
+ return tag
+ normalized = tag.strip()
+ if not normalized:
+ return None
+ return normalized.lower()
+
+
+def get_parser():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ parser.add_argument(
+ "--epoch",
+ type=int,
+ default=30,
+ help="""It specifies the checkpoint to use for decoding.
+ Note: Epoch counts from 1.
+ You can specify --avg to use more checkpoints for model averaging.""",
+ )
+
+ parser.add_argument(
+ "--iter",
+ type=int,
+ default=0,
+ help="""If positive, --epoch is ignored and it
+ will use the checkpoint exp_dir/checkpoint-iter.pt.
+ You can specify --avg to use more checkpoints for model averaging.
+ """,
+ )
+
+ parser.add_argument(
+ "--avg",
+ type=int,
+ default=15,
+ help="Number of checkpoints to average. Automatically select "
+ "consecutive checkpoints before the checkpoint specified by "
+ "'--epoch' and '--iter'",
+ )
+
+ parser.add_argument(
+ "--use-averaged-model",
+ type=str2bool,
+ default=True,
+ help="Whether to load averaged model. Currently it only supports "
+ "using --epoch. If True, it would decode with the averaged model "
+ "over the epoch range from `epoch-avg` (excluded) to `epoch`."
+ "Actually only the models with epoch number of `epoch-avg` and "
+ "`epoch` are loaded for averaging. ",
+ )
+ parser.add_argument(
+ "--model-name",
+ type=str,
+ default=None,
+ help="Specify a model name",
+ )
+ parser.add_argument(
+ "--exp-dir",
+ type=str,
+ default="zipformer/exp",
+ help="The experiment dir",
+ )
+
+ parser.add_argument(
+ "--bpe-model",
+ type=str,
+ default="data/lang_bpe_500/bpe.model",
+ help="Path to the BPE model",
+ )
+
+ parser.add_argument(
+ "--lang-dir",
+ type=Path,
+ default="data/lang_bpe_500",
+ help="The lang dir containing word table and LG graph",
+ )
+
+ parser.add_argument(
+ "--decoding-method",
+ type=str,
+ default="greedy_search",
+ help="""Possible values are:
+ - greedy_search
+ - beam_search
+ - modified_beam_search
+ - modified_beam_search_LODR
+ - fast_beam_search
+ - fast_beam_search_nbest
+ - fast_beam_search_nbest_oracle
+ - fast_beam_search_nbest_LG
+ If you use fast_beam_search_nbest_LG, you have to specify
+ `--lang-dir`, which should contain `LG.pt`.
+ """,
+ )
+
+ parser.add_argument(
+ "--decoding-method-dir",
+ type=str,
+ default="modified_beam_search",
+ )
+ parser.add_argument(
+ "--beam-size",
+ type=int,
+ default=4,
+ help="""An integer indicating how many candidates we will keep for each
+ frame. Used only when --decoding-method is beam_search or
+ modified_beam_search.""",
+ )
+
+ parser.add_argument(
+ "--beam",
+ type=float,
+ default=20.0,
+ help="""A floating point value to calculate the cutoff score during beam
+ search (i.e., `cutoff = max-score - beam`), which is the same as the
+ `beam` in Kaldi.
+ Used only when --decoding-method is fast_beam_search,
+ fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle
+ """,
+ )
+
+ parser.add_argument(
+ "--ngram-lm-scale",
+ type=float,
+ default=0.01,
+ help="""
+ Used only when --decoding-method is fast_beam_search_nbest_LG.
+ It specifies the scale for n-gram LM scores.
+ """,
+ )
+
+ parser.add_argument(
+ "--max-contexts",
+ type=int,
+ default=8,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--max-states",
+ type=int,
+ default=64,
+ help="""Used only when --decoding-method is
+ fast_beam_search, fast_beam_search_nbest, fast_beam_search_nbest_LG,
+ and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="The context size in the decoder. 1 means bigram; " "2 means tri-gram",
+ )
+ parser.add_argument(
+ "--max-sym-per-frame",
+ type=int,
+ default=1,
+ help="""Maximum number of symbols per frame.
+ Used only when --decoding-method is greedy_search""",
+ )
+
+ parser.add_argument(
+ "--num-paths",
+ type=int,
+ default=200,
+ help="""Number of paths for nbest decoding.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--nbest-scale",
+ type=float,
+ default=0.5,
+ help="""Scale applied to lattice scores when computing nbest paths.
+ Used only when the decoding method is fast_beam_search_nbest,
+ fast_beam_search_nbest_LG, and fast_beam_search_nbest_oracle""",
+ )
+
+ parser.add_argument(
+ "--use-shallow-fusion",
+ type=str2bool,
+ default=False,
+ help="""Use neural network LM for shallow fusion.
+ If you want to use LODR, you will also need to set this to true
+ """,
+ )
+
+ parser.add_argument(
+ "--lm-type",
+ type=str,
+ default="rnn",
+ help="Type of NN lm",
+ choices=["rnn", "transformer"],
+ )
+
+ parser.add_argument(
+ "--lm-scale-shallow-fusion",
+ type=float,
+ default=0.3,
+ help="""The scale of the neural network LM
+ Used only when `--use-shallow-fusion` is set to True.
+ """,
+ )
+
+ parser.add_argument(
+ "--tokens-ngram",
+ type=int,
+ default=2,
+ help="""The order of the ngram lm.
+ """,
+ )
+
+ parser.add_argument(
+ "--backoff-id",
+ type=int,
+ default=500,
+ help="ID of the backoff symbol in the ngram LM",
+ )
+
+ parser.add_argument(
+ "--context-score",
+ type=float,
+ default=2,
+ help="""
+ The bonus score of each token for the context biasing words/phrases.
+ Used only when --decoding-method is modified_beam_search and
+ modified_beam_search_LODR.
+ """,
+ )
+
+ parser.add_argument(
+ "--context-file",
+ type=str,
+ default="",
+ help="""
+ The path of the context biasing lists, one word/phrase each line
+ Used only when --decoding-method is modified_beam_search and
+ modified_beam_search_LODR.
+ """,
+ )
+ parser.add_argument(
+ "--dump-moe-routing-stats",
+ type=str2bool,
+ default=False,
+ )
+
+ parser.add_argument(
+ "--skip-scoring",
+ type=str2bool,
+ default=False,
+ )
+
+ parser.add_argument(
+ "--compute-cer",
+ type=str2bool,
+ default=False,
+ help="If True, compute character error rate.",
+ )
+
+ parser.add_argument(
+ "--remove-punctuation",
+ type=str2bool,
+ default=False,
+ help="If True, remove punctuation symbols.",
+ )
+
+ parser.add_argument("--asr-decode", type=str2bool, default=False)
+ parser.add_argument("--ast-decode", type=str2bool, default=False)
+
+ # --- in get_parser() ---
+
+ parser.add_argument(
+ "--blank-penalty-asr",
+ type=float,
+ default=0.0,
+ )
+
+ parser.add_argument(
+ "--blank-penalty-st",
+ type=float,
+ default=0.0,
+ )
+
+ parser.add_argument("--use-tgt", type=str2bool, default=False)
+ parser.add_argument(
+ "--lang-tgt",
+ type=str,
+ default="",
+ )
+ parser.add_argument(
+ "--force-first-lang",
+ type=str2bool,
+ default=False,
+ )
+
+ add_model_arguments(parser)
+
+ return parser
+
+
+class _OutputLinearWithBlankPenalty(nn.Module):
+ def __init__(self, linear: nn.Module, blank_id: int, penalty: float):
+ super().__init__()
+ self.linear = linear
+ self.blank_id = int(blank_id)
+ self.penalty = float(penalty)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ z = self.linear(x) # logits
+ if self.penalty > 0.0:
+ z[..., self.blank_id] = z[..., self.blank_id] - self.penalty
+ return z
+
+ @property
+ def weight(self):
+ return self.linear.weight
+
+ @property
+ def bias(self):
+ return self.linear.bias
+
+
+class JoinerWithBlankPenalty(nn.Module):
+ def __init__(self, joiner: nn.Module, blank_id: int, penalty: float):
+ super().__init__()
+ self.inner = joiner
+ self.blank_id = int(blank_id)
+ self.penalty = float(penalty)
+
+ self.output_linear = _OutputLinearWithBlankPenalty(
+ getattr(joiner, "output_linear"), self.blank_id, self.penalty
+ )
+
+ def forward(self, *args, **kwargs):
+ z = self.inner(*args, **kwargs)
+ if self.penalty > 0.0:
+ z[..., self.blank_id] = z[..., self.blank_id] - self.penalty
+ return z
+
+ def __getattr__(self, name: str):
+ try:
+ return super().__getattr__(name)
+ except AttributeError:
+ return getattr(self.inner, name)
+
+
+def decode_one_batch(
+ params: AttributeDict,
+ model: nn.Module,
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ batch: dict,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+ context_graph: Optional[ContextGraph] = None,
+ LM: Optional[LmScorer] = None,
+ ngram_lm=None,
+ ngram_lm_scale: float = 0.0,
+ srt_lang_ids: Optional[torch.Tensor] = None,
+ tgt_lang_ids: Optional[torch.Tensor] = None,
+) -> Dict[str, List[List[str]]]:
+ """Decode one batch and return the result in a dict. The dict has the
+ following format:
+
+ - key: It indicates the setting used for decoding. For example,
+ if greedy_search is used, it would be "greedy_search"
+ If beam search with a beam size of 7 is used, it would be
+ "beam_7"
+ - value: It contains the decoding result. `len(value)` equals to
+ batch size. `value[i]` is the decoding result for the i-th
+ utterance in the given batch.
+ Args:
+ params:
+ It's the return value of :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ batch:
+ It is the return value from iterating
+ `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+ for the format of the `batch`.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding-method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ LM:
+ A neural network language model.
+ ngram_lm:
+ A ngram language model
+ ngram_lm_scale:
+ The scale for the ngram language model.
+ Returns:
+ Return the decoding result. See above description for the format of
+ the returned dict.
+ """
+ device = next(model.parameters()).device
+ feature = batch["inputs"]
+ assert feature.ndim == 3
+
+ feature = feature.to(device)
+ # at entry, feature is (N, T, C)
+
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ if params.causal:
+ # this seems to cause insertions at the end of the utterance if used with zipformer.
+ pad_len = 30
+ feature_lens += pad_len
+ feature = torch.nn.functional.pad(
+ feature,
+ pad=(0, 0, 0, pad_len),
+ value=LOG_EPS,
+ )
+
+ collect_moe_stats = bool(getattr(params, "dump_moe_routing_stats", False))
+ if collect_moe_stats:
+ (
+ asr_encoder_out,
+ asr_encoder_out_lens,
+ st_encoder_out,
+ st_encoder_out_lens,
+ moe_loss,
+ moe_weights_asr,
+ moe_weights_st,
+ ) = model.forward_encoder(
+ feature,
+ feature_lens,
+ srt_lang_ids,
+ tgt_lang_ids,
+ enable_st=params.enable_st,
+ return_moe_weights=True,
+ )
+ else:
+ (
+ asr_encoder_out,
+ asr_encoder_out_lens,
+ st_encoder_out,
+ st_encoder_out_lens,
+ moe_loss,
+ ) = model.forward_encoder(
+ feature,
+ feature_lens,
+ srt_lang_ids,
+ tgt_lang_ids,
+ enable_st=params.enable_st,
+ )
+ moe_weights_asr = None
+ moe_weights_st = None
+
+ asr_hyps = []
+ st_hyps = []
+ if params.decoding_method == "modified_beam_search":
+ # ===== ASR =====
+ if params.asr_decode:
+ joiner_asr = model.joiner_asr
+ if getattr(params, "blank_penalty_asr", 0.0) > 0.0:
+ joiner_asr = JoinerWithBlankPenalty(
+ joiner=model.joiner_asr,
+ blank_id=params.blank_id_asr,
+ penalty=params.blank_penalty_asr,
+ )
+
+ asr_hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=asr_encoder_out,
+ encoder_out_lens=asr_encoder_out_lens,
+ decoder=model.decoder_asr,
+ joiner=joiner_asr,
+ beam=params.beam_size,
+ context_graph=context_graph,
+ )
+ for asr_hyp in sp_asr.decode(asr_hyp_tokens):
+ asr_hyps.append(asr_hyp.split())
+
+ # ===== ST =====
+ if params.ast_decode:
+ joiner_st = model.joiner_st
+ if getattr(params, "blank_penalty_st", 0.0) > 0.0:
+ joiner_st = JoinerWithBlankPenalty(
+ joiner=model.joiner_st,
+ blank_id=params.blank_id_st,
+ penalty=params.blank_penalty_st,
+ )
+
+ lang_tgt = sp_st.piece_to_id(params.lang_tgt)
+ st_hyp_tokens = modified_beam_search(
+ model=model,
+ encoder_out=st_encoder_out,
+ encoder_out_lens=st_encoder_out_lens,
+ decoder=model.decoder_st,
+ joiner=joiner_st,
+ beam=params.beam_size,
+ context_graph=context_graph,
+ lang_token_id=lang_tgt,
+ force_first_lang=params.force_first_lang,
+ )
+ for st_hyp in sp_st.decode(st_hyp_tokens):
+ st_hyps.append(st_hyp.split())
+
+ else:
+ raise NotImplementedError(
+ f"Decoding method '{params.decoding_method}' is not supported "
+ "in this dual ASR/ST decoder. Use 'modified_beam_search'."
+ )
+
+ asr_prefix = f"asr_{params.decoding_method}"
+ st_prefix = f"st_{params.decoding_method}"
+ asr_result: Dict[str, List[List[str]]] = dict()
+ st_result: Dict[str, List[List[str]]] = dict()
+
+ if "modified_beam_search" in params.decoding_method:
+ asr_prefix += f"_beam-size-{params.beam_size}"
+ st_prefix += f"_beam-size-{params.beam_size}"
+ if params.has_contexts:
+ asr_prefix += f"_context-score-{params.context_score}"
+ asr_result = {asr_prefix: asr_hyps}
+ st_result = {st_prefix: st_hyps}
+ else:
+ raise NotImplementedError(
+ f"Decoding method '{params.decoding_method}' is not supported."
+ )
+
+ moe_batch_stats = None
+ if collect_moe_stats:
+ moe_batch_stats = dict()
+ if moe_weights_asr is not None and srt_lang_ids is not None:
+ moe_batch_stats["asr"] = (
+ srt_lang_ids.detach().cpu(),
+ moe_weights_asr.mean(dim=0).detach().cpu(),
+ )
+ if moe_weights_st is not None and tgt_lang_ids is not None:
+ moe_batch_stats["st"] = (
+ tgt_lang_ids.detach().cpu(),
+ moe_weights_st.mean(dim=0).detach().cpu(),
+ )
+ if not moe_batch_stats:
+ moe_batch_stats = None
+
+ return asr_result, st_result, moe_batch_stats
+
+
+def _extract_st_texts_and_lang_ids(
+ supervisions: List[Dict[str, Any]],
+ use_tgt: bool,
+ tgt_lang2id: Dict[str, int],
+ default_lang: str = None,
+):
+ default_lang = _normalize_lang_tag(default_lang)
+ st_texts = []
+ lang_ids: list[int] = []
+
+ for cut in supervisions["cut"]:
+ for supervision in cut.supervisions:
+ if hasattr(supervision, "custom") and "st_text" in supervision.custom:
+ lang_tag = None
+ if "lang" in supervision.custom and supervision.custom["lang"]:
+ lang_tag = _normalize_lang_tag(supervision.custom["lang"])
+ if lang_tag is None and default_lang is not None:
+ lang_tag = default_lang
+
+ if use_tgt:
+ if lang_tag is None:
+ raise ValueError(
+ "Missing custom['lang'] and no default_lang provided."
+ )
+ if lang_tag not in tgt_lang2id:
+ raise KeyError(
+ f"Unknown target language tag: {lang_tag}. "
+ f"Known: {list(tgt_lang2id.keys())}"
+ )
+ supervision.custom["st_text"] = (
+ f"<2{lang_tag}>" + supervision.custom["st_text"]
+ )
+ lang_ids.append(tgt_lang2id[lang_tag])
+ else:
+ if lang_tag is None:
+ lang_ids.append(0)
+ else:
+ lang_ids.append(tgt_lang2id.get(lang_tag, 0))
+
+ st_texts.append(supervision.custom["st_text"])
+
+ tgt_lang_ids = torch.tensor(lang_ids, dtype=torch.long)
+ return st_texts, tgt_lang_ids
+
+
+def asr_source_lang_tensor(
+ supervisions: Dict[str, Any],
+ srt_lang2id: Dict[str, int],
+ *,
+ strict: bool = True,
+) -> torch.LongTensor:
+ tags: List[Optional[str]] = []
+
+ cuts: Iterable[Any] = supervisions.get("cut", [])
+ for cut in cuts:
+ sups = getattr(cut, "supervisions", None)
+ if sups is None and isinstance(cut, dict):
+ sups = cut.get("supervisions", [])
+ if not sups:
+ continue
+
+ for sup in sups:
+ if isinstance(sup, dict):
+ text = sup.get("text")
+ lang = sup.get("language")
+ else:
+ text = getattr(sup, "text", None)
+ lang = getattr(sup, "language", None)
+
+ if text is None:
+ continue
+ if lang == "English":
+ lang = "en"
+ lang = _normalize_lang_tag(lang)
+
+ if lang is None:
+ raise KeyError("Missing supervision['language'] for an ASR sample.")
+ tags.append(lang)
+
+ if "text" in supervisions and isinstance(supervisions["text"], list):
+ assert len(tags) == len(
+ supervisions["text"]
+ ), f"The number of ASR languages ({len(tags)}) is inconsistent with the number of texts ({len(supervisions['text'])})."
+
+ if strict:
+ ids = []
+ for t in tags:
+ if t not in srt_lang2id:
+ raise ValueError(
+ f"Unknown source language: {t}. Known: {list(srt_lang2id.keys())}"
+ )
+ ids.append(srt_lang2id[t])
+ else:
+ ids = [srt_lang2id.get(t, 0) for t in tags]
+
+ return torch.tensor(ids, dtype=torch.long)
+
+
+def decode_dataset(
+ dl: torch.utils.data.DataLoader,
+ params: AttributeDict,
+ model: nn.Module,
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ word_table: Optional[k2.SymbolTable] = None,
+ decoding_graph: Optional[k2.Fsa] = None,
+ context_graph: Optional[ContextGraph] = None,
+ LM: Optional[LmScorer] = None,
+ ngram_lm=None,
+ ngram_lm_scale: float = 0.0,
+) -> Dict[str, List[Tuple[str, List[str], List[str]]]]:
+ """Decode dataset.
+
+ Args:
+ dl:
+ PyTorch's dataloader containing the dataset to decode.
+ params:
+ It is returned by :func:`get_params`.
+ model:
+ The neural model.
+ sp:
+ The BPE model.
+ word_table:
+ The word symbol table.
+ decoding_graph:
+ The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+ only when --decoding-method is fast_beam_search, fast_beam_search_nbest,
+ fast_beam_search_nbest_oracle, and fast_beam_search_nbest_LG.
+ Returns:
+ Return a dict, whose key may be "greedy_search" if greedy search
+ is used, or it may be "beam_7" if beam size of 7 is used.
+ Its value is a list of tuples. Each tuple contains two elements:
+ The first is the reference transcript, and the second is the
+ predicted result.
+ """
+ num_cuts = 0
+
+ try:
+ num_batches = len(dl)
+ except TypeError:
+ num_batches = "?"
+
+ if params.decoding_method == "greedy_search":
+ log_interval = 50
+ else:
+ log_interval = 20
+
+ results_asr = defaultdict(list)
+ results_st = defaultdict(list)
+
+ collect_moe_stats = bool(getattr(params, "dump_moe_routing_stats", False))
+ # NOTE: In this repo the actual MoE modules are attached as `asr_moe_layer` / `ast_moe_layer`
+ # (see `zipformer/train_cross_node_jsrt.py:get_model()`).
+ base_model = model.module if hasattr(model, "module") else model
+ num_asr_experts = getattr(
+ getattr(base_model, "asr_moe_layer", None), "num_experts", 0
+ )
+ num_st_experts = getattr(
+ getattr(base_model, "ast_moe_layer", None), "num_experts", 0
+ )
+ moe_stats_asr = (
+ defaultdict(lambda: torch.zeros(num_asr_experts, dtype=torch.float64))
+ if collect_moe_stats and num_asr_experts > 0
+ else None
+ )
+ moe_counts_asr = defaultdict(int) if moe_stats_asr is not None else None
+ moe_stats_st = (
+ defaultdict(lambda: torch.zeros(num_st_experts, dtype=torch.float64))
+ if collect_moe_stats and num_st_experts > 0
+ else None
+ )
+ moe_counts_st = defaultdict(int) if moe_stats_st is not None else None
+
+ def _accumulate_moe_stats(storage, counts, batch_info):
+ if storage is None or counts is None or batch_info is None:
+ return
+ lang_ids, weights = batch_info
+ if lang_ids is None or weights is None:
+ return
+ lang_list = lang_ids.tolist()
+ for idx, w in zip(lang_list, weights):
+ storage[idx] += w.to(storage[idx].dtype)
+ counts[idx] += 1
+
+ def _log_moe_stats(task_name, storage, counts, lang_list):
+ if storage is None or counts is None:
+ return
+ logging.info("===== MoE routing stats (%s) =====", task_name)
+ for lang_id, total in sorted(storage.items()):
+ count = counts[lang_id]
+ if count == 0:
+ continue
+ avg = (total / count).tolist()
+ dist = ", ".join(f"e{i}:{val:.3f}" for i, val in enumerate(avg))
+ lang = lang_list[lang_id] if 0 <= lang_id < len(lang_list) else str(lang_id)
+ logging.info(" %s (id=%d, n=%d): %s", lang, lang_id, count, dist)
+
+ for batch_idx, batch in enumerate(dl):
+ supervisions = batch["supervisions"]
+
+ srt_lang_ids = None
+ if params.asr_decode:
+ texts_asr: List[str] = supervisions["text"]
+ # Backward compatibility:
+ # - old flag: asr_moe_use_src_embed
+ # - new flag (train_cross_node_jsrt.py): asr_src
+ use_asr_src = bool(
+ getattr(params, "asr_src", False)
+ or getattr(params, "asr_moe_use_src_embed", False)
+ )
+ srt_lang_ids = (
+ asr_source_lang_tensor(supervisions, params.srt_lang2id, strict=True)
+ if use_asr_src
+ else None
+ )
+
+ tgt_lang_ids = None
+ if params.enable_st and getattr(params, "ast_tgt", True):
+ texts_st, tgt_lang_ids = _extract_st_texts_and_lang_ids(
+ supervisions, params.use_tgt, params.tgt_lang2id
+ )
+ # Optional composite ids (only if these legacy flags exist)
+ if (
+ getattr(params, "use_srctgt_lang_ids", False)
+ and not getattr(params, "ast_use_src_tgt_embed", False)
+ and srt_lang_ids is not None
+ ):
+ tgt_lang_ids = srt_lang_ids * params.num_tgt_langs_ast + tgt_lang_ids
+ if getattr(params, "use_no_lang_ids", False):
+ tgt_lang_ids = None
+ else:
+ texts_st, tgt_lang_ids = [], None
+
+ cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+ decode_outputs = decode_one_batch(
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ decoding_graph=decoding_graph,
+ context_graph=context_graph,
+ word_table=word_table,
+ batch=batch,
+ LM=LM,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ srt_lang_ids=srt_lang_ids,
+ tgt_lang_ids=tgt_lang_ids,
+ )
+ if isinstance(decode_outputs, tuple) and len(decode_outputs) == 3:
+ asr_hyps_dict, st_hyps_dict, batch_moe_stats = decode_outputs
+ else:
+ asr_hyps_dict, st_hyps_dict = decode_outputs
+ batch_moe_stats = None
+
+ if collect_moe_stats and batch_moe_stats:
+ _accumulate_moe_stats(
+ moe_stats_asr,
+ moe_counts_asr,
+ batch_moe_stats.get("asr"),
+ )
+ _accumulate_moe_stats(
+ moe_stats_st,
+ moe_counts_st,
+ batch_moe_stats.get("st"),
+ )
+
+ if params.asr_decode:
+ for name, hyps in asr_hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts_asr)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts_asr):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+ results_asr[name].extend(this_batch)
+ if params.ast_decode:
+ for name, hyps in st_hyps_dict.items():
+ this_batch = []
+ assert len(hyps) == len(texts_st)
+ for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts_st):
+ ref_words = ref_text.split()
+ this_batch.append((cut_id, ref_words, hyp_words))
+
+ results_st[name].extend(this_batch)
+
+ num_cuts += len(cut_ids)
+
+ if batch_idx % log_interval == 0:
+ batch_str = f"{batch_idx}/{num_batches}"
+ logging.info(f"batch {batch_str}, cuts processed until now is {num_cuts}")
+ if collect_moe_stats:
+ _log_moe_stats("ASR", moe_stats_asr, moe_counts_asr, params.srt_lang_list)
+ st_lang_labels = (
+ getattr(params, "srctgt_lang_list", params.tgt_lang_list)
+ if getattr(params, "use_srctgt_lang_ids", False)
+ else params.tgt_lang_list
+ )
+ _log_moe_stats("ST", moe_stats_st, moe_counts_st, st_lang_labels)
+
+ return results_asr, results_st
+
+
+def save_asr_output(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict_asr: Dict[str, List[Tuple[str, List[str], List[str]]]],
+ results_dict_st: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+ """
+ Save text produced by ASR.
+ """
+ if params.asr_decode:
+ for key, results in results_dict_asr.items():
+
+ recogs_filename = (
+ params.res_dir / f"recogs-asr-{test_set_name}-{params.suffix}.txt"
+ )
+
+ results = sorted(results)
+ store_transcripts(filename=recogs_filename, texts=results)
+
+ logging.info(f"The transcripts are stored in {recogs_filename}")
+
+ if params.ast_decode:
+ for key, results in results_dict_st.items():
+
+ recogs_filename = (
+ params.res_dir / f"recogs-st-{test_set_name}-{params.suffix}.txt"
+ )
+
+ results = sorted(results)
+ store_transcripts(filename=recogs_filename, texts=results)
+
+ logging.info(f"The transcripts are stored in {recogs_filename}")
+
+
+def asr_save_wer_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str], Tuple]]],
+):
+ """
+ Save WER and per-utterance word alignments.
+ """
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ # The following prints out WERs, per-word error statistics and aligned
+ # ref/hyp pairs.
+ errs_filename = params.res_dir / f"errs-asr-{test_set_name}-{params.suffix}.txt"
+ with open(errs_filename, "w", encoding="utf8") as fd:
+ wer = write_error_stats(
+ # fd, f"{test_set_name}-{key}", results, enable_log=True
+ fd,
+ f"{test_set_name}-{key}",
+ results,
+ enable_log=True,
+ compute_CER=params.compute_cer,
+ remove_punctuation=params.remove_punctuation,
+ )
+ test_set_wers[key] = wer
+
+ logging.info(f"Wrote detailed error stats to {errs_filename}")
+
+ test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
+
+ wer_filename = (
+ params.res_dir / f"wer-asr-summary-{test_set_name}-{params.suffix}.txt"
+ )
+
+ with open(wer_filename, "w", encoding="utf8") as fd:
+ print("settings\tWER", file=fd)
+ for key, val in test_set_wers:
+ print(f"{key}\t{val}", file=fd)
+
+ s = f"\nFor {test_set_name}, WER of different settings are:\n"
+ note = f"\tbest for {test_set_name}"
+ for key, val in test_set_wers:
+ s += f"{key}\t{val}{note}\n"
+ note = ""
+ logging.info(s)
+
+
+def st_save_wer_results(
+ params: AttributeDict,
+ test_set_name: str,
+ results_dict: Dict[str, List[Tuple[str, List[str], List[str], Tuple]]],
+):
+ """
+ Save WER and per-utterance word alignments.
+ """
+ test_set_wers = dict()
+ for key, results in results_dict.items():
+ # The following prints out WERs, per-word error statistics and aligned
+ # ref/hyp pairs.
+ errs_filename = params.res_dir / f"errs-st-{test_set_name}-{params.suffix}.txt"
+ with open(errs_filename, "w", encoding="utf8") as fd:
+ wer = write_error_stats(
+ # fd, f"{test_set_name}-{key}", results, enable_log=True
+ fd,
+ f"{test_set_name}-{key}",
+ results,
+ enable_log=True,
+ compute_CER=params.compute_cer,
+ remove_punctuation=params.remove_punctuation,
+ )
+ test_set_wers[key] = wer
+
+ logging.info(f"Wrote detailed error stats to {errs_filename}")
+
+ test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
+
+ wer_filename = (
+ params.res_dir / f"wer-st-summary-{test_set_name}-{params.suffix}.txt"
+ )
+
+ with open(wer_filename, "w", encoding="utf8") as fd:
+ print("settings\tWER", file=fd)
+ for key, val in test_set_wers:
+ print(f"{key}\t{val}", file=fd)
+
+ s = f"\nFor {test_set_name}, WER of different settings are:\n"
+ note = f"\tbest for {test_set_name}"
+ for key, val in test_set_wers:
+ s += f"{key}\t{val}{note}\n"
+ note = ""
+ logging.info(s)
+
+
+@torch.no_grad()
+def main():
+ parser = get_parser()
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ LmScorer.add_arguments(parser)
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ params = get_params()
+ params.update(vars(args))
+
+ # enable AudioCache
+ set_caching_enabled(True) # lhotse
+
+ assert params.decoding_method in (
+ "greedy_search",
+ "beam_search",
+ "fast_beam_search",
+ "fast_beam_search_nbest",
+ "fast_beam_search_nbest_LG",
+ "fast_beam_search_nbest_oracle",
+ "modified_beam_search",
+ "modified_beam_search_LODR",
+ "modified_beam_search_lm_shallow_fusion",
+ "modified_beam_search_lm_rescore",
+ "modified_beam_search_lm_rescore_LODR",
+ )
+ params.res_dir = params.exp_dir / params.decoding_method_dir
+
+ if os.path.exists(params.context_file):
+ params.has_contexts = True
+ else:
+ params.has_contexts = False
+
+ if params.iter > 0:
+ params.suffix = f"iter-{params.iter}_avg-{params.avg}"
+ else:
+ params.suffix = f"epoch-{params.epoch}_avg-{params.avg}"
+
+ if params.causal:
+ assert (
+ "," not in params.chunk_size
+ ), "chunk_size should be one value in decoding."
+ assert (
+ "," not in params.left_context_frames
+ ), "left_context_frames should be one value in decoding."
+ params.suffix += f"_chunk-{params.chunk_size}"
+ params.suffix += f"_left-context-{params.left_context_frames}"
+
+ if "fast_beam_search" in params.decoding_method:
+ params.suffix += f"_beam-{params.beam}"
+ params.suffix += f"_max-contexts-{params.max_contexts}"
+ params.suffix += f"_max-states-{params.max_states}"
+ if "nbest" in params.decoding_method:
+ params.suffix += f"_nbest-scale-{params.nbest_scale}"
+ params.suffix += f"_num-paths-{params.num_paths}"
+ if "LG" in params.decoding_method:
+ params.suffix += f"_ngram-lm-scale-{params.ngram_lm_scale}"
+ elif "beam_search" in params.decoding_method:
+ params.suffix += f"__{params.decoding_method}__beam-size-{params.beam_size}"
+ if params.decoding_method in (
+ "modified_beam_search",
+ "modified_beam_search_LODR",
+ ):
+ if params.has_contexts:
+ params.suffix += f"-context-score-{params.context_score}"
+ else:
+ params.suffix += f"_context-{params.context_size}"
+ params.suffix += f"_max-sym-per-frame-{params.max_sym_per_frame}"
+
+ if params.use_shallow_fusion:
+ params.suffix += f"_{params.lm_type}-lm-scale-{params.lm_scale_shallow_fusion}"
+ if "LODR" in params.decoding_method:
+ params.suffix += (
+ f"_LODR-{params.tokens_ngram}gram-scale-{params.ngram_lm_scale}"
+ )
+
+ if params.use_averaged_model:
+ params.suffix += "_use-averaged-model"
+
+ setup_logger(f"{params.res_dir}/log-decode-{params.suffix}")
+ logging.info("Decoding started")
+
+ device = torch.device("cpu")
+ if torch.cuda.is_available():
+ device = torch.device("cuda", 0)
+
+ logging.info(f"Device: {device}")
+
+ # Tokenizers
+ sp_asr = spm.SentencePieceProcessor()
+ sp_asr.load(params.bpe_model_asr)
+ sp_st = spm.SentencePieceProcessor()
+ sp_st.load(params.bpe_model_st)
+
+ # Ids and vocab sizes per task
+ params.blank_id_asr = sp_asr.piece_to_id("")
+ params.sos_id_asr = params.eos_id_asr = sp_asr.piece_to_id("")
+ params.vocab_size_asr = sp_asr.get_piece_size()
+
+ params.blank_id_st = (
+ sp_st.piece_to_id("") if sp_st.piece_to_id("") != -1 else 0
+ )
+ params.sos_id_st = params.eos_id_st = (
+ sp_st.piece_to_id("") if sp_st.piece_to_id("") != -1 else 1
+ )
+ params.vocab_size_st = sp_st.get_piece_size()
+
+ params.tgt_lang_list = [s.strip() for s in params.tgt_langs.split(",") if s.strip()]
+ params.tgt_lang2id = {lg: i for i, lg in enumerate(params.tgt_lang_list)}
+ params.num_tgt_langs_ast = len(params.tgt_lang_list)
+
+ params.srt_lang_list = [s.strip() for s in params.srt_langs.split(",") if s.strip()]
+ params.srt_lang2id = {lg: i for i, lg in enumerate(params.srt_lang_list)}
+ params.num_srt_langs_asr = len(params.srt_lang_list)
+ params.srctgt_lang_list = build_srctgt_lang_list(
+ params.srt_lang_list, params.tgt_lang_list
+ )
+
+ logging.info(params)
+
+ logging.info("About to create model")
+ model = get_model(params)
+
+ if not params.use_averaged_model:
+ if params.model_name:
+ load_checkpoint(f"{params.exp_dir}/{params.model_name}", model)
+ elif params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ elif params.avg == 1:
+ load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
+ else:
+ start = params.epoch - params.avg + 1
+ filenames = []
+ for i in range(start, params.epoch + 1):
+ if i >= 1:
+ filenames.append(f"{params.exp_dir}/epoch-{i}.pt")
+ logging.info(f"averaging {filenames}")
+ model.to(device)
+ model.load_state_dict(average_checkpoints(filenames, device=device))
+ else:
+ if params.iter > 0:
+ filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[
+ : params.avg + 1
+ ]
+ if len(filenames) == 0:
+ raise ValueError(
+ f"No checkpoints found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ elif len(filenames) < params.avg + 1:
+ raise ValueError(
+ f"Not enough checkpoints ({len(filenames)}) found for"
+ f" --iter {params.iter}, --avg {params.avg}"
+ )
+ filename_start = filenames[-1]
+ filename_end = filenames[0]
+ logging.info(
+ "Calculating the averaged model over iteration checkpoints"
+ f" from {filename_start} (excluded) to {filename_end}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+ else:
+ assert params.avg > 0, params.avg
+ start = params.epoch - params.avg
+ assert start >= 1, start
+ filename_start = f"{params.exp_dir}/epoch-{start}.pt"
+ filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt"
+ logging.info(
+ f"Calculating the averaged model over epoch range from "
+ f"{start} (excluded) to {params.epoch}"
+ )
+ model.to(device)
+ model.load_state_dict(
+ average_checkpoints_with_averaged_model(
+ filename_start=filename_start,
+ filename_end=filename_end,
+ device=device,
+ )
+ )
+
+ model.to(device)
+ model.eval()
+
+ # only load the neural network LM if required
+ if params.use_shallow_fusion or params.decoding_method in (
+ "modified_beam_search_lm_rescore",
+ "modified_beam_search_lm_rescore_LODR",
+ "modified_beam_search_lm_shallow_fusion",
+ "modified_beam_search_LODR",
+ ):
+ LM = LmScorer(
+ lm_type=params.lm_type,
+ params=params,
+ device=device,
+ lm_scale=params.lm_scale_shallow_fusion,
+ )
+ LM.to(device)
+ LM.eval()
+ else:
+ LM = None
+
+ # only load N-gram LM when needed
+ if params.decoding_method == "modified_beam_search_lm_rescore_LODR":
+ try:
+ import kenlm
+ except ImportError:
+ print("Please install kenlm first. You can use")
+ print(" pip install https://github.com/kpu/kenlm/archive/master.zip")
+ print("to install it")
+ import sys
+
+ sys.exit(-1)
+ ngram_file_name = str(params.lang_dir / f"{params.tokens_ngram}gram.arpa")
+ logging.info(f"lm filename: {ngram_file_name}")
+ ngram_lm = kenlm.Model(ngram_file_name)
+ ngram_lm_scale = None # use a list to search
+
+ elif params.decoding_method == "modified_beam_search_LODR":
+ lm_filename = f"{params.tokens_ngram}gram.fst.txt"
+ logging.info(f"Loading token level lm: {lm_filename}")
+ ngram_lm = NgramLm(
+ str(params.lang_dir / lm_filename),
+ backoff_id=params.backoff_id,
+ is_binary=False,
+ )
+ logging.info(f"num states: {ngram_lm.lm.num_states}")
+ ngram_lm_scale = params.ngram_lm_scale
+ else:
+ ngram_lm = None
+ ngram_lm_scale = None
+
+ if "fast_beam_search" in params.decoding_method:
+ if params.decoding_method == "fast_beam_search_nbest_LG":
+ lexicon = Lexicon(params.lang_dir)
+ word_table = lexicon.word_table
+ lg_filename = params.lang_dir / "LG.pt"
+ logging.info(f"Loading {lg_filename}")
+ decoding_graph = k2.Fsa.from_dict(
+ torch.load(lg_filename, map_location=device, weights_only=False)
+ )
+ decoding_graph.scores *= params.ngram_lm_scale
+ else:
+ word_table = None
+ decoding_graph = k2.trivial_graph(params.vocab_size_asr - 1, device=device)
+ else:
+ decoding_graph = None
+ word_table = None
+
+ if "modified_beam_search" in params.decoding_method:
+ if os.path.exists(params.context_file):
+ contexts = []
+ with open(params.context_file) as f:
+ for line in f:
+ contexts.append((sp_asr.encode(line.strip()), 0.0))
+ context_graph = ContextGraph(params.context_score)
+ context_graph.build(contexts)
+ else:
+ context_graph = None
+ else:
+ context_graph = None
+
+ num_param = sum([p.numel() for p in model.parameters()])
+ logging.info(f"Number of model parameters: {num_param}")
+
+ args.return_cuts = True
+ librispeech = LibriSpeechAsrDataModule(args)
+ test_cuts = librispeech.test_cuts()
+
+ def remove_short_and_long_utt(c: Cut) -> bool:
+ if c.duration < 0.3 or c.duration > 30:
+ return False
+ return True
+
+ test_cuts = test_cuts.filter(remove_short_and_long_utt)
+
+ test_dl = librispeech.test_dataloaders(test_cuts)
+
+ name = "test"
+ results_dict_asr, results_dict_st, = decode_dataset(
+ dl=test_dl,
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ word_table=word_table,
+ decoding_graph=decoding_graph,
+ context_graph=context_graph,
+ LM=LM,
+ ngram_lm=ngram_lm,
+ ngram_lm_scale=ngram_lm_scale,
+ )
+
+ save_asr_output(
+ params=params,
+ test_set_name=name,
+ results_dict_asr=results_dict_asr,
+ results_dict_st=results_dict_st,
+ )
+
+ if not params.skip_scoring:
+ asr_save_wer_results(
+ params=params,
+ test_set_name=name,
+ results_dict=results_dict_asr,
+ )
+ # if params.ast_use_asr_data:
+ # st_save_wer_results(
+ # params=params,
+ # test_set_name=name,
+ # results_dict=results_dict_st,
+ # )
+
+ logging.info("Done!")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc.sh
new file mode 100644
index 0000000000..c634bdfbe6
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc.sh
@@ -0,0 +1,131 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 0 \
+ --use-ctc-st 0 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --force-first-lang 0 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --enable-st 0 \
+ --dump-moe-routing-stats 0 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
+
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_moe.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_moe.sh
new file mode 100644
index 0000000000..b2d1dd1389
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_moe.sh
@@ -0,0 +1,131 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 0 \
+ --use-ctc-st 0 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --force-first-lang 0 \
+ --asr-moe 1 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 8 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --enable-st 0 \
+ --dump-moe-routing-stats 0 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
+
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_s_bias.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_s_bias.sh
new file mode 100644
index 0000000000..241dbb1bb4
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_s_bias.sh
@@ -0,0 +1,130 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 0 \
+ --use-ctc-st 0 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --force-first-lang 0 \
+ --asr-moe 0 \
+ --asr-src 1 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --enable-st 0 \
+ --dump-moe-routing-stats 0 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_sc_moe.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_sc_moe.sh
new file mode 100644
index 0000000000..7830c3719a
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage1/decode_cr_ctc_sc_moe.sh
@@ -0,0 +1,130 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 0 \
+ --use-ctc-st 0 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --force-first-lang 0 \
+ --asr-moe 1 \
+ --asr-src 1 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 8 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --enable-st 0 \
+ --dump-moe-routing-stats 1 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2m.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2m.sh
new file mode 100644
index 0000000000..8388ba2f93
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2m.sh
@@ -0,0 +1,194 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/es_de/europarl.es_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_de/europarl.fr_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_de/europarl.it_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_de/europarl.nl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_de/europarl.pl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_de/europarl.pt_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_de/europarl.ro_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_de/europarl.en_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_es/europarl.de_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_es/europarl.fr_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_es/europarl.it_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_es/europarl.nl_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_es/europarl.pl_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_es/europarl.pt_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_es/europarl.ro_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_es/europarl.en_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_fr/europarl.de_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_fr/europarl.es_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_fr/europarl.it_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_fr/europarl.nl_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_fr/europarl.pl_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_fr/europarl.pt_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_fr/europarl.ro_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_it/europarl.de_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_it/europarl.fr_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_it/europarl.es_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_it/europarl.nl_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_it/europarl.pl_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_it/europarl.pt_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_it/europarl.ro_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_it/europarl.en_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_nl/europarl.de_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_nl/europarl.fr_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_nl/europarl.it_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_nl/europarl.es_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_nl/europarl.pl_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_nl/europarl.pt_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_nl/europarl.ro_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_nl/europarl.en_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_pl/europarl.de_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_pl/europarl.fr_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_pl/europarl.it_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_pl/europarl.nl_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_pl/europarl.es_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_pl/europarl.pt_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_pl/europarl.ro_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_pl/europarl.en_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_pt/europarl.de_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_pt/europarl.fr_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_pt/europarl.it_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_pt/europarl.nl_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_pt/europarl.pl_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_pt/europarl.es_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_pt/europarl.ro_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_pt/europarl.en_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_ro/europarl.de_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_ro/europarl.fr_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_ro/europarl.it_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_ro/europarl.nl_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_ro/europarl.pl_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_ro/europarl.pt_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_ro/europarl.es_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_ro/europarl.en_ro.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+ tgt_lang_token="<2${tgt_lang}>"
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 1 \
+ --use-ctc-st 1 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --lang-tgt "${tgt_lang_token}" \
+ --force-first-lang 1 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --dump-moe-routing-stats 0 \
+ --enable-st 1 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2o.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2o.sh
new file mode 100644
index 0000000000..2dfc3c690d
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_hent_srt_m2o.sh
@@ -0,0 +1,131 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast_de/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl/x_de"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/es_de/europarl.es_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_de/europarl.fr_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_de/europarl.it_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_de/europarl.nl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_de/europarl.pl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_de/europarl.pt_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_de/europarl.ro_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_de/europarl.en_de.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+ tgt_lang_token="<2${tgt_lang}>"
+
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 1 \
+ --use-ctc-st 1 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --lang-tgt "${tgt_lang_token}" \
+ --force-first-lang 1 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --dump-moe-routing-stats 0 \
+ --enable-st 1 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
diff --git a/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_lcma_srt.sh b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_lcma_srt.sh
new file mode 100644
index 0000000000..611492ea15
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decode/stage2/decode_lcma_srt.sh
@@ -0,0 +1,194 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Decoding script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+steps=0
+epochs=0
+avg=1
+MODEL_NAME="best-valid-loss.pt"
+EXP_DIR="exp/europarl"
+
+TEST_CUTS_PATHS=(
+ data/Europarl-ST/manifests/es_de/europarl.es_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_de/europarl.fr_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_de/europarl.it_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_de/europarl.nl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_de/europarl.pl_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_de/europarl.pt_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_de/europarl.ro_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_de/europarl.en_de.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_en/europarl.de_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_en/europarl.fr_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_en/europarl.it_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_en/europarl.nl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_en/europarl.pl_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_en/europarl.pt_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_en/europarl.ro_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_en/europarl.es_en.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_es/europarl.de_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_es/europarl.fr_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_es/europarl.it_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_es/europarl.nl_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_es/europarl.pl_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_es/europarl.pt_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_es/europarl.ro_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_es/europarl.en_es.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_fr/europarl.de_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_fr/europarl.es_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_fr/europarl.it_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_fr/europarl.nl_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_fr/europarl.pl_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_fr/europarl.pt_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_fr/europarl.ro_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_fr/europarl.en_fr.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_it/europarl.de_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_it/europarl.fr_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_it/europarl.es_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_it/europarl.nl_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_it/europarl.pl_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_it/europarl.pt_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_it/europarl.ro_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_it/europarl.en_it.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_nl/europarl.de_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_nl/europarl.fr_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_nl/europarl.it_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_nl/europarl.es_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_nl/europarl.pl_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_nl/europarl.pt_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_nl/europarl.ro_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_nl/europarl.en_nl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_pl/europarl.de_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_pl/europarl.fr_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_pl/europarl.it_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_pl/europarl.nl_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_pl/europarl.es_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_pl/europarl.pt_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_pl/europarl.ro_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_pl/europarl.en_pl.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_pt/europarl.de_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_pt/europarl.fr_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_pt/europarl.it_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_pt/europarl.nl_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_pt/europarl.pl_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_pt/europarl.es_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/ro_pt/europarl.ro_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_pt/europarl.en_pt.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/de_ro/europarl.de_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/fr_ro/europarl.fr_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/it_ro/europarl.it_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/nl_ro/europarl.nl_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pl_ro/europarl.pl_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/pt_ro/europarl.pt_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/es_ro/europarl.es_ro.test_cuts.jsonl.gz
+ data/Europarl-ST/manifests/en_ro/europarl.en_ro.test_cuts.jsonl.gz
+)
+
+compute_cer=False
+
+
+manifest_dir="data/Europarl-ST/manifests"
+decoding_method="modified_beam_search"
+beam_size=20
+
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+for TEST_CUTS_PATH in "${TEST_CUTS_PATHS[@]}"; do
+ lang_pair_dir=$(basename "$(dirname "$TEST_CUTS_PATH")")
+ src_lang=${lang_pair_dir%%_*}
+ tgt_lang=${lang_pair_dir##*_}
+ if [[ "$src_lang" == "$lang_pair_dir" ]]; then
+ tgt_lang=$src_lang
+ fi
+
+ decoding_method_dir="modified_beam_search_beam20_${src_lang}_to_${tgt_lang}_cuts_test"
+ compute_cer_current=$compute_cer
+ tgt_lang_token="<2${tgt_lang}>"
+ python lcma_srt/decode.py \
+ --iter $steps \
+ --avg $avg \
+ --use-averaged-model 0 \
+ --exp-dir $EXP_DIR \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --manifest-dir $manifest_dir \
+ --decoding-method $decoding_method \
+ --beam-size $beam_size \
+ --max-duration 500 \
+ --compute-cer $compute_cer_current \
+ --remove-punctuation True \
+ --causal 0 \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --chunk-size -1 \
+ --test-name $TEST_CUTS_PATH \
+ --left-context-frames -1 \
+ --use-ctc-asr 1 \
+ --asr-decode 1 \
+ --ast-decode 1 \
+ --use-ctc-st 1 \
+ --blank-penalty-st 2.0 \
+ --decoding-method-dir $decoding_method_dir \
+ --num-heads-asr ${num_heads_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --lang-tgt "${tgt_lang_token}" \
+ --force-first-lang 1 \
+ --asr-moe 1 \
+ --asr-src 1 \
+ --ast-moe 1 \
+ --ast-tgt 1 \
+ --num-experts-asr 8 \
+ --num-experts-ast 16 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.015 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --dump-moe-routing-stats 1 \
+ --enable-st 1 \
+ --model-name $MODEL_NAME
+done
+
+echo "--- Decoding script finished at $(date) ---"
diff --git a/egs/europarl_st/SRT/lcma_srt/decoder.py b/egs/europarl_st/SRT/lcma_srt/decoder.py
new file mode 100644
index 0000000000..e9b3ef8ca0
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/decoder.py
@@ -0,0 +1,132 @@
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from scaling import Balancer
+
+
+class Decoder(nn.Module):
+ """This class modifies the stateless decoder from the following paper:
+
+ RNN-transducer with stateless prediction network
+ https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=9054419
+
+ It removes the recurrent connection from the decoder, i.e., the prediction
+ network. Different from the above paper, it adds an extra Conv1d
+ right after the embedding layer.
+
+ TODO: Implement https://arxiv.org/pdf/2109.07513.pdf
+ """
+
+ def __init__(
+ self,
+ vocab_size: int,
+ decoder_dim: int,
+ blank_id: int,
+ context_size: int,
+ ):
+ """
+ Args:
+ vocab_size:
+ Number of tokens of the modeling unit including blank.
+ decoder_dim:
+ Dimension of the input embedding, and of the decoder output.
+ blank_id:
+ The ID of the blank symbol.
+ context_size:
+ Number of previous words to use to predict the next word.
+ 1 means bigram; 2 means trigram. n means (n+1)-gram.
+ """
+ super().__init__()
+
+ self.embedding = nn.Embedding(
+ num_embeddings=vocab_size,
+ embedding_dim=decoder_dim,
+ )
+ # the balancers are to avoid any drift in the magnitude of the
+ # embeddings, which would interact badly with parameter averaging.
+ self.balancer = Balancer(
+ decoder_dim,
+ channel_dim=-1,
+ min_positive=0.0,
+ max_positive=1.0,
+ min_abs=0.5,
+ max_abs=1.0,
+ prob=0.05,
+ )
+
+ self.blank_id = blank_id
+
+ assert context_size >= 1, context_size
+ self.context_size = context_size
+ self.vocab_size = vocab_size
+
+ if context_size > 1:
+ self.conv = nn.Conv1d(
+ in_channels=decoder_dim,
+ out_channels=decoder_dim,
+ kernel_size=context_size,
+ padding=0,
+ groups=decoder_dim // 4, # group size == 4
+ bias=False,
+ )
+ self.balancer2 = Balancer(
+ decoder_dim,
+ channel_dim=-1,
+ min_positive=0.0,
+ max_positive=1.0,
+ min_abs=0.5,
+ max_abs=1.0,
+ prob=0.05,
+ )
+ else:
+ # To avoid `RuntimeError: Module 'Decoder' has no attribute 'conv'`
+ # when inference with torch.jit.script and context_size == 1
+ self.conv = nn.Identity()
+ self.balancer2 = nn.Identity()
+
+ def forward(self, y: torch.Tensor, need_pad: bool = True) -> torch.Tensor:
+ """
+ Args:
+ y:
+ A 2-D tensor of shape (N, U).
+ need_pad:
+ True to left pad the input. Should be True during training.
+ False to not pad the input. Should be False during inference.
+ Returns:
+ Return a tensor of shape (N, U, decoder_dim).
+ """
+ y = y.to(torch.int64)
+ # this stuff about clamp() is a temporary fix for a mismatch
+ # at utterance start, we use negative ids in beam_search.py
+ embedding_out = self.embedding(y.clamp(min=0)) * (y >= 0).unsqueeze(-1)
+
+ embedding_out = self.balancer(embedding_out)
+
+ if self.context_size > 1:
+ embedding_out = embedding_out.permute(0, 2, 1)
+ if need_pad is True:
+ embedding_out = F.pad(embedding_out, pad=(self.context_size - 1, 0))
+ else:
+ # During inference time, there is no need to do extra padding
+ # as we only need one output
+ assert embedding_out.size(-1) == self.context_size
+ embedding_out = self.conv(embedding_out)
+ embedding_out = embedding_out.permute(0, 2, 1)
+ embedding_out = F.relu(embedding_out)
+ embedding_out = self.balancer2(embedding_out)
+
+ return embedding_out
diff --git a/egs/europarl_st/SRT/lcma_srt/encoder_interface.py b/egs/europarl_st/SRT/lcma_srt/encoder_interface.py
new file mode 100644
index 0000000000..257facce4f
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/encoder_interface.py
@@ -0,0 +1,43 @@
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+
+
+class EncoderInterface(nn.Module):
+ def forward(
+ self, x: torch.Tensor, x_lens: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ x:
+ A tensor of shape (batch_size, input_seq_len, num_features)
+ containing the input features.
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames
+ in `x` before padding.
+ Returns:
+ Return a tuple containing two tensors:
+ - encoder_out, a tensor of (batch_size, out_seq_len, output_dim)
+ containing unnormalized probabilities, i.e., the output of a
+ linear layer.
+ - encoder_out_lens, a tensor of shape (batch_size,) containing
+ the number of frames in `encoder_out` before padding.
+ """
+ raise NotImplementedError("Please implement it in a subclass")
diff --git a/egs/europarl_st/SRT/lcma_srt/joiner.py b/egs/europarl_st/SRT/lcma_srt/joiner.py
new file mode 100644
index 0000000000..0406efe834
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/joiner.py
@@ -0,0 +1,67 @@
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import torch.nn as nn
+from scaling import ScaledLinear
+
+
+class Joiner(nn.Module):
+ def __init__(
+ self,
+ encoder_dim: int,
+ decoder_dim: int,
+ joiner_dim: int,
+ vocab_size: int,
+ ):
+ super().__init__()
+
+ self.encoder_proj = ScaledLinear(encoder_dim, joiner_dim, initial_scale=0.25)
+ self.decoder_proj = ScaledLinear(decoder_dim, joiner_dim, initial_scale=0.25)
+ self.output_linear = nn.Linear(joiner_dim, vocab_size)
+
+ def forward(
+ self,
+ encoder_out: torch.Tensor,
+ decoder_out: torch.Tensor,
+ project_input: bool = True,
+ ) -> torch.Tensor:
+ """
+ Args:
+ encoder_out:
+ Output from the encoder. Its shape is (N, T, s_range, C).
+ decoder_out:
+ Output from the decoder. Its shape is (N, T, s_range, C).
+ project_input:
+ If true, apply input projections encoder_proj and decoder_proj.
+ If this is false, it is the user's responsibility to do this
+ manually.
+ Returns:
+ Return a tensor of shape (N, T, s_range, C).
+ """
+ assert encoder_out.ndim == decoder_out.ndim, (
+ encoder_out.shape,
+ decoder_out.shape,
+ )
+
+ if project_input:
+ logit = self.encoder_proj(encoder_out) + self.decoder_proj(decoder_out)
+ else:
+ logit = encoder_out + decoder_out
+
+ logit = self.output_linear(torch.tanh(logit))
+
+ return logit
diff --git a/egs/europarl_st/SRT/lcma_srt/label_smoothing.py b/egs/europarl_st/SRT/lcma_srt/label_smoothing.py
new file mode 100644
index 0000000000..52d2eda3bb
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/label_smoothing.py
@@ -0,0 +1,109 @@
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+
+class LabelSmoothingLoss(torch.nn.Module):
+ """
+ Implement the LabelSmoothingLoss proposed in the following paper
+ https://arxiv.org/pdf/1512.00567.pdf
+ (Rethinking the Inception Architecture for Computer Vision)
+
+ """
+
+ def __init__(
+ self,
+ ignore_index: int = -1,
+ label_smoothing: float = 0.1,
+ reduction: str = "sum",
+ ) -> None:
+ """
+ Args:
+ ignore_index:
+ ignored class id
+ label_smoothing:
+ smoothing rate (0.0 means the conventional cross entropy loss)
+ reduction:
+ It has the same meaning as the reduction in
+ `torch.nn.CrossEntropyLoss`. It can be one of the following three
+ values: (1) "none": No reduction will be applied. (2) "mean": the
+ mean of the output is taken. (3) "sum": the output will be summed.
+ """
+ super().__init__()
+ assert 0.0 <= label_smoothing < 1.0, f"{label_smoothing}"
+ assert reduction in ("none", "sum", "mean"), reduction
+ self.ignore_index = ignore_index
+ self.label_smoothing = label_smoothing
+ self.reduction = reduction
+
+ def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+ """
+ Compute loss between x and target.
+
+ Args:
+ x:
+ prediction of dimension
+ (batch_size, input_length, number_of_classes).
+ target:
+ target masked with self.ignore_index of
+ dimension (batch_size, input_length).
+
+ Returns:
+ A scalar tensor containing the loss without normalization.
+ """
+ assert x.ndim == 3
+ assert target.ndim == 2
+ assert x.shape[:2] == target.shape
+ num_classes = x.size(-1)
+ x = x.reshape(-1, num_classes)
+ # Now x is of shape (N*T, C)
+
+ # We don't want to change target in-place below,
+ # so we make a copy of it here
+ target = target.clone().reshape(-1)
+
+ ignored = target == self.ignore_index
+
+ # See https://github.com/k2-fsa/icefall/issues/240
+ # and https://github.com/k2-fsa/icefall/issues/297
+ # for why we don't use target[ignored] = 0 here
+ target = torch.where(ignored, torch.zeros_like(target), target)
+
+ true_dist = torch.nn.functional.one_hot(target, num_classes=num_classes).to(x)
+
+ true_dist = (
+ true_dist * (1 - self.label_smoothing) + self.label_smoothing / num_classes
+ )
+
+ # Set the value of ignored indexes to 0
+ #
+ # See https://github.com/k2-fsa/icefall/issues/240
+ # and https://github.com/k2-fsa/icefall/issues/297
+ # for why we don't use true_dist[ignored] = 0 here
+ true_dist = torch.where(
+ ignored.unsqueeze(1).repeat(1, true_dist.shape[1]),
+ torch.zeros_like(true_dist),
+ true_dist,
+ )
+
+ loss = -1 * (torch.log_softmax(x, dim=1) * true_dist)
+ if self.reduction == "sum":
+ return loss.sum()
+ elif self.reduction == "mean":
+ return loss.sum() / (~ignored).sum()
+ else:
+ return loss.sum(dim=-1)
diff --git a/egs/europarl_st/SRT/lcma_srt/model.py b/egs/europarl_st/SRT/lcma_srt/model.py
new file mode 100644
index 0000000000..476c4c6d5b
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/model.py
@@ -0,0 +1,793 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import List, Optional, Tuple
+
+import k2
+import torch
+import torch.nn as nn
+from encoder_interface import EncoderInterface
+from lhotse.dataset import SpecAugment
+from moe_adapter import MoEAdapterDense
+from scaling import ScaledLinear
+from torch import Tensor, nn
+
+from icefall.utils import add_sos, make_pad_mask, time_warp
+
+
+class LCMASRTModel(nn.Module):
+ def __init__(
+ self,
+ encoder_embed: nn.Module,
+ enc_asr: EncoderInterface,
+ enc_st: EncoderInterface,
+ decoder_asr: Optional[nn.Module] = None,
+ joiner_asr: Optional[nn.Module] = None,
+ attention_decoder_asr: Optional[nn.Module] = None,
+ decoder_st: Optional[nn.Module] = None,
+ joiner_st: Optional[nn.Module] = None,
+ attention_decoder_st: Optional[nn.Module] = None,
+ encoder_dim_asr: int = 384,
+ encoder_dim_st: int = 384,
+ decoder_dim_asr: int = 512,
+ decoder_dim_st: int = 512,
+ vocab_size_asr: int = 500,
+ vocab_size_st: int = 6000,
+ output_downsampling_factor_asr: int = 2,
+ output_downsampling_factor_st: int = 2,
+ num_srt_langs_asr: int = 0,
+ num_tgt_langs_ast: int = 0,
+ num_experts_asr: int = 4,
+ num_experts_ast: int = 8,
+ entropy_reg_asr: float = 0.0,
+ entropy_reg_ast: float = 0.0,
+ temperature_asr: float = 1.0,
+ temperature_ast: float = 1.0,
+ asr_moe: bool = True,
+ asr_src: bool = True,
+ ast_moe: bool = True,
+ ast_tgt: bool = True,
+ # task flags
+ use_transducer: bool = True,
+ use_ctc_asr: bool = True,
+ use_ctc_st: bool = True,
+ # CR-CTC / Attn flags (CR-CTC is supported; attn optional)
+ use_attention_decoder: bool = False,
+ freeze_asr: bool = False,
+ freeze_frontend: bool = False,
+ ):
+ super().__init__()
+ assert (
+ use_transducer or use_ctc_asr
+ ), f"At least one of them should be True, but got use_transducer={use_transducer}, use_ctc={use_ctc_asr}"
+ assert isinstance(enc_asr, EncoderInterface)
+ assert isinstance(enc_st, EncoderInterface)
+ self.is_srt = True
+ self.encoder_embed = encoder_embed
+ self.enc_asr = enc_asr
+ self.enc_st = enc_st
+
+ self.use_transducer = use_transducer
+ self.use_ctc_asr = use_ctc_asr
+ self.use_ctc_st = use_ctc_st
+ self.use_attention_decoder = use_attention_decoder
+
+ self.output_downsampling_factor_asr = output_downsampling_factor_asr
+ self.output_downsampling_factor_st = output_downsampling_factor_st
+ # ------ ASR head ------
+ if self.use_transducer:
+ assert decoder_asr is not None and joiner_asr is not None
+ assert hasattr(decoder_asr, "blank_id")
+ self.decoder_asr = decoder_asr
+ self.joiner_asr = joiner_asr
+ self.simple_am_proj_asr = ScaledLinear(
+ encoder_dim_asr, vocab_size_asr, initial_scale=0.25
+ )
+ self.simple_lm_proj_asr = ScaledLinear(
+ decoder_dim_asr, vocab_size_asr, initial_scale=0.25
+ )
+ else:
+ assert decoder_asr is None
+ assert joiner_asr is None
+
+ if self.use_ctc_asr:
+ self.ctc_asr = nn.Sequential(
+ nn.Dropout(p=0.1),
+ nn.Linear(encoder_dim_asr, vocab_size_asr),
+ nn.LogSoftmax(dim=-1),
+ )
+
+ # ------ ST head ------
+ if self.use_transducer:
+ assert decoder_st is not None and joiner_st is not None
+ assert hasattr(decoder_st, "blank_id")
+ self.decoder_st = decoder_st
+ self.joiner_st = joiner_st
+ self.simple_am_proj_st = ScaledLinear(
+ encoder_dim_st, vocab_size_st, initial_scale=0.25
+ )
+ self.simple_lm_proj_st = ScaledLinear(
+ decoder_dim_st, vocab_size_st, initial_scale=0.25
+ )
+ else:
+ assert decoder_st is None
+ assert joiner_st is None
+ if self.use_ctc_st:
+ self.ctc_st = nn.Sequential(
+ nn.Dropout(p=0.1),
+ nn.Linear(encoder_dim_st, vocab_size_st),
+ nn.LogSoftmax(dim=-1),
+ )
+
+ # optional attention decoder heads could be added similarly if needed
+ if use_attention_decoder:
+ self.attention_decoder_asr = attention_decoder_asr
+ self.attention_decoder_st = attention_decoder_st
+ else:
+ assert attention_decoder_asr is None
+ assert attention_decoder_st is None
+
+ self.freeze_asr = freeze_asr
+ self.freeze_frontend = freeze_frontend
+ if self.freeze_asr:
+ self._apply_freeze_asr()
+
+ self.asr_moe = asr_moe
+ self.asr_src = asr_src
+ self.ast_moe = ast_moe
+ self.ast_tgt = ast_tgt
+
+ self.asr_moe_layer: Optional[MoEAdapterDense] = None
+ if self.asr_moe:
+ num_langs_asr = num_srt_langs_asr if self.asr_src else 0
+ self.asr_moe_layer = MoEAdapterDense(
+ d_model=enc_asr.output_dim,
+ num_experts=num_experts_asr,
+ hidden_mult=1.3,
+ num_tasks=0,
+ num_langs=num_langs_asr,
+ dropout=0.1,
+ entropy_reg=entropy_reg_asr,
+ temperature=temperature_asr,
+ )
+ self.lang_embed_asr = (
+ nn.Embedding(num_srt_langs_asr, enc_asr.output_dim)
+ if (not self.asr_moe and self.asr_src)
+ else None
+ )
+
+ self.ast_moe_layer: Optional[MoEAdapterDense] = None
+ if self.ast_moe:
+ num_langs_ast = num_tgt_langs_ast if self.ast_tgt else 0
+ self.ast_moe_layer = MoEAdapterDense(
+ d_model=enc_st.output_dim,
+ num_experts=num_experts_ast,
+ hidden_mult=1.3,
+ num_tasks=0,
+ num_langs=num_langs_ast,
+ num_src_langs=0,
+ num_tgt_langs=0,
+ dropout=0.1,
+ entropy_reg=entropy_reg_ast,
+ temperature=temperature_ast,
+ )
+ self.lang_embed_ast = (
+ nn.Embedding(num_tgt_langs_ast, enc_st.output_dim)
+ if (not self.ast_moe and self.ast_tgt)
+ else None
+ )
+
+ def _apply_freeze_asr(self):
+ to_freeze = []
+ if self.freeze_frontend and hasattr(self, "encoder_embed"):
+ to_freeze += [self.encoder_embed]
+ to_freeze += [self.enc_asr]
+ for name in [
+ "decoder_asr",
+ "joiner_asr",
+ "ctc_asr",
+ "attention_decoder_asr",
+ "simple_am_proj_asr",
+ "simple_lm_proj_asr",
+ ]:
+ if hasattr(self, name) and getattr(self, name) is not None:
+ to_freeze += [getattr(self, name)]
+
+ for m in to_freeze:
+ m.eval()
+ for p in m.parameters():
+ p.requires_grad = False
+
+ def output_downsampling(self, x_lengths):
+ # class Downsample has this rounding behavior..
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ lengths = (x_lengths + 1) // 2
+ else:
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ lengths = (x_lengths + 1) // 2
+
+ return lengths
+
+ def forward_encoder(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ srt_lang_ids: Optional[torch.Tensor] = None,
+ tgt_lang_ids: Optional[torch.Tensor] = None,
+ enable_st: bool = True,
+ return_moe_weights: bool = False,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+ """Compute encoder outputs.
+ Args:
+ x:
+ A 3-D tensor of shape (N, T, C).
+ x_lens:
+ A 1-D tensor of shape (N,). It contains the number of frames in `x`
+ before padding.
+ """
+
+ if self.freeze_frontend:
+ with torch.no_grad():
+ x, x_lens = self.encoder_embed(x, x_lens)
+ else:
+ x, x_lens = self.encoder_embed(x, x_lens)
+
+ pad0 = make_pad_mask(x_lens)
+ x_tnc = x.permute(1, 0, 2) # (T,N,C)
+
+ if self.freeze_asr:
+ with torch.no_grad():
+ asr_x, asr_len = self.enc_asr(x_tnc, x_lens, pad0)
+ else:
+ asr_x, asr_len = self.enc_asr(x_tnc, x_lens, pad0) # (T1,N,C1), (N,)
+
+ if self.freeze_asr:
+ asr_x = asr_x.detach()
+
+ w_asr: Optional[torch.Tensor] = None
+ if self.asr_moe and self.asr_moe_layer is not None:
+ if self.asr_moe_layer.lang_embed is not None:
+ if srt_lang_ids is None:
+ raise ValueError(
+ "srt_lang_ids is required when asr_moe=True and asr_src=True"
+ )
+ lang_ids = srt_lang_ids.to(dtype=torch.long, device=asr_x.device)
+ else:
+ lang_ids = None
+ asr_x_for_asr, w_asr = self.asr_moe_layer(asr_x, lang_ids=lang_ids)
+ elif self.asr_src and self.lang_embed_asr is not None:
+ if srt_lang_ids is None:
+ raise ValueError(
+ "srt_lang_ids is required when asr_src=True and asr_moe=False"
+ )
+ srt_lang_ids = srt_lang_ids.to(dtype=torch.long, device=asr_x.device)
+ lang_bias = (
+ self.lang_embed_asr(srt_lang_ids).unsqueeze(0).to(dtype=asr_x.dtype)
+ )
+ asr_x_for_asr, w_asr = asr_x + lang_bias, None
+ else:
+ asr_x_for_asr, w_asr = asr_x, None
+
+ asr_tnc = self.enc_asr.downsample_output(asr_x_for_asr)
+ asr_output_len = self.output_downsampling(asr_len)
+ asr_ntc = asr_tnc.permute(1, 0, 2)
+
+ moe_terms: List[torch.Tensor] = []
+ if self.asr_moe and w_asr is not None and self.asr_moe_layer is not None:
+ pad_tb_asr = make_pad_mask(asr_len).transpose(0, 1)
+ moe_ent_asr = self.asr_moe_layer.router_entropy_loss(
+ w=w_asr, pad_mask_tb=pad_tb_asr
+ )
+ moe_terms.append(moe_ent_asr)
+
+ w_ast: Optional[torch.Tensor] = None
+ if not enable_st:
+ dummy_st = torch.empty(0, device=asr_ntc.device)
+ dummy_lens = torch.zeros_like(asr_output_len)
+ if moe_terms:
+ moe_ent_loss = torch.stack(moe_terms).mean()
+ else:
+ moe_ent_loss = torch.zeros((), device=asr_ntc.device)
+ if return_moe_weights:
+ return (
+ asr_ntc,
+ asr_output_len,
+ dummy_st,
+ dummy_lens,
+ moe_ent_loss,
+ w_asr,
+ w_ast,
+ )
+ return asr_ntc, asr_output_len, dummy_st, dummy_lens, moe_ent_loss
+
+ from scaling import convert_num_channels
+
+ st_in = convert_num_channels(asr_x_for_asr, self.enc_st.encoder_dim[0])
+ st_x, st_len = self.enc_st(st_in, asr_len, make_pad_mask(asr_len))
+
+ if self.ast_moe and self.ast_moe_layer is not None:
+ if self.ast_moe_layer.lang_embed is not None:
+ if tgt_lang_ids is None:
+ raise ValueError(
+ "tgt_lang_ids is required when ast_moe=True and ast_tgt=True"
+ )
+ lang_ids = tgt_lang_ids.to(dtype=torch.long, device=st_x.device)
+ else:
+ lang_ids = None
+ ast_x_for_ast, w_ast = self.ast_moe_layer(
+ st_x,
+ lang_ids=lang_ids,
+ )
+ elif self.ast_tgt and self.lang_embed_ast is not None:
+ if tgt_lang_ids is None:
+ raise ValueError(
+ "tgt_lang_ids is required when ast_tgt=True and ast_moe=False"
+ )
+ tgt_lang_ids = tgt_lang_ids.to(dtype=torch.long, device=st_x.device)
+ lang_bias = (
+ self.lang_embed_ast(tgt_lang_ids).unsqueeze(0).to(dtype=st_x.dtype)
+ )
+ ast_x_for_ast, w_ast = st_x + lang_bias, None
+ else:
+ ast_x_for_ast, w_ast = st_x, None
+
+ if self.output_downsampling_factor_st == 2:
+ st_tnc = self.enc_st.downsample_output(ast_x_for_ast)
+ st_lens = self.output_downsampling(st_len)
+ elif self.output_downsampling_factor_st == 1:
+ st_tnc = ast_x_for_ast
+ st_lens = st_len
+
+ st_ntc = st_tnc.permute(1, 0, 2) # (N,T2,C2)
+
+ assert torch.all(asr_output_len > 0) and torch.all(st_lens > 0), (
+ asr_output_len,
+ st_lens,
+ )
+ # return asr_ntc, asr_output_len, st_ntc, st_lens
+
+ if self.ast_moe and w_ast is not None and self.ast_moe_layer is not None:
+ pad_tb_st = make_pad_mask(st_len).transpose(0, 1)
+ moe_ent_st = self.ast_moe_layer.router_entropy_loss(
+ w=w_ast, pad_mask_tb=pad_tb_st
+ )
+ moe_terms.append(moe_ent_st)
+
+ if moe_terms:
+ moe_ent_loss = torch.stack(moe_terms).mean()
+ else:
+ moe_ent_loss = torch.zeros((), device=asr_ntc.device)
+
+ if return_moe_weights:
+ return (
+ asr_ntc,
+ asr_output_len,
+ st_ntc,
+ st_lens,
+ moe_ent_loss,
+ w_asr,
+ w_ast,
+ )
+
+ return asr_ntc, asr_output_len, st_ntc, st_lens, moe_ent_loss
+
+ def forward_transducer(
+ self,
+ encoder_out: torch.Tensor, # (N, T, C)
+ encoder_out_lens: torch.Tensor, # (N,)
+ y: k2.RaggedTensor,
+ y_lens: torch.Tensor,
+ decoder: nn.Module,
+ joiner: nn.Module,
+ simple_lm_proj: nn.Module,
+ simple_am_proj: nn.Module,
+ prune_range: int = 5,
+ am_scale: float = 0.0,
+ lm_scale: float = 0.0,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute Transducer loss.
+ Args:
+ encoder_out:
+ Encoder output, of shape (N, T, C).
+ encoder_out_lens:
+ Encoder output lengths, of shape (N,).
+ y:
+ A ragged tensor with 2 axes [utt][label]. It contains labels of each
+ utterance.
+ prune_range:
+ The prune range for rnnt loss, it means how many symbols(context)
+ we are considering for each frame to compute the loss.
+ am_scale:
+ The scale to smooth the loss with am (output of encoder network)
+ part
+ lm_scale:
+ The scale to smooth the loss with lm (output of predictor network)
+ part
+ """
+ # Now for the decoder, i.e., the prediction network
+ blank_id = decoder.blank_id
+ sos_y = add_sos(y, sos_id=blank_id)
+
+ # sos_y_padded: [B, S + 1], start with SOS.
+ sos_y_padded = sos_y.pad(mode="constant", padding_value=blank_id)
+
+ # decoder_out: [B, S + 1, decoder_dim]
+ decoder_out = decoder(sos_y_padded)
+
+ # Note: y does not start with SOS
+ # y_padded : [B, S]
+ y_padded = y.pad(mode="constant", padding_value=blank_id)
+
+ y_padded = y_padded.to(torch.int64)
+ boundary = torch.zeros(
+ (encoder_out.size(0), 4),
+ dtype=torch.int64,
+ device=encoder_out.device,
+ )
+ boundary[:, 2] = y_lens
+ boundary[:, 3] = encoder_out_lens
+
+ lm = simple_lm_proj(decoder_out)
+ am = simple_am_proj(encoder_out)
+
+ # if self.training and random.random() < 0.25:
+ # lm = penalize_abs_values_gt(lm, 100.0, 1.0e-04)
+ # if self.training and random.random() < 0.25:
+ # am = penalize_abs_values_gt(am, 30.0, 1.0e-04)
+
+ with torch.cuda.amp.autocast(enabled=False):
+ simple_loss, (px_grad, py_grad) = k2.rnnt_loss_smoothed(
+ lm=lm.float(),
+ am=am.float(),
+ symbols=y_padded,
+ termination_symbol=blank_id,
+ lm_only_scale=lm_scale,
+ am_only_scale=am_scale,
+ boundary=boundary,
+ reduction="sum",
+ return_grad=True,
+ )
+
+ # ranges : [B, T, prune_range]
+ ranges = k2.get_rnnt_prune_ranges(
+ px_grad=px_grad,
+ py_grad=py_grad,
+ boundary=boundary,
+ s_range=prune_range,
+ )
+
+ # am_pruned : [B, T, prune_range, encoder_dim]
+ # lm_pruned : [B, T, prune_range, decoder_dim]
+ am_pruned, lm_pruned = k2.do_rnnt_pruning(
+ am=joiner.encoder_proj(encoder_out),
+ lm=joiner.decoder_proj(decoder_out),
+ ranges=ranges,
+ )
+
+ # logits : [B, T, prune_range, vocab_size]
+
+ # project_input=False since we applied the decoder's input projections
+ # prior to do_rnnt_pruning (this is an optimization for speed).
+ logits = joiner(am_pruned, lm_pruned, project_input=False)
+
+ with torch.cuda.amp.autocast(enabled=False):
+ pruned_loss = k2.rnnt_loss_pruned(
+ logits=logits.float(),
+ symbols=y_padded,
+ ranges=ranges,
+ termination_symbol=blank_id,
+ boundary=boundary,
+ reduction="sum",
+ )
+
+ return simple_loss, pruned_loss
+
+ def forward_ctc(
+ self,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ targets: torch.Tensor,
+ target_lengths: torch.Tensor,
+ ctc_output: nn.Module,
+ ) -> torch.Tensor:
+ """Compute CTC loss.
+ Args:
+ encoder_out:
+ Encoder output, of shape (N, T, C).
+ encoder_out_lens:
+ Encoder output lengths, of shape (N,).
+ targets:
+ Target Tensor of shape (sum(target_lengths)). The targets are assumed
+ to be un-padded and concatenated within 1 dimension.
+ """
+ # Compute CTC log-prob
+ ctc_output = ctc_output(encoder_out) # (N, T, C)
+ ctc_loss = torch.nn.functional.ctc_loss(
+ log_probs=ctc_output.permute(1, 0, 2), # (T, N, C)
+ targets=targets.cpu(),
+ input_lengths=encoder_out_lens.cpu(),
+ target_lengths=target_lengths.cpu(),
+ reduction="sum",
+ )
+ return ctc_loss
+
+ def forward_cr_ctc(
+ self,
+ encoder_out: torch.Tensor,
+ encoder_out_lens: torch.Tensor,
+ targets: torch.Tensor,
+ target_lengths: torch.Tensor,
+ ctc_output: nn.Module,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute CTC loss with consistency regularization loss.
+ Args:
+ encoder_out:
+ Encoder output, of shape (2 * N, T, C).
+ encoder_out_lens:
+ Encoder output lengths, of shape (2 * N,).
+ targets:
+ Target Tensor of shape (2 * sum(target_lengths)). The targets are assumed
+ to be un-padded and concatenated within 1 dimension.
+ """
+ # Compute CTC loss
+ ctc_output = ctc_output(encoder_out) # (2 * N, T, C)
+ ctc_loss = torch.nn.functional.ctc_loss(
+ log_probs=ctc_output.permute(1, 0, 2), # (T, 2 * N, C)
+ targets=targets.cpu(),
+ input_lengths=encoder_out_lens.cpu(),
+ target_lengths=target_lengths.cpu(),
+ reduction="sum",
+ )
+
+ # Compute consistency regularization loss
+ exchanged_targets = ctc_output.detach().chunk(2, dim=0)
+ exchanged_targets = torch.cat(
+ [exchanged_targets[1], exchanged_targets[0]], dim=0
+ ) # exchange: [x1, x2] -> [x2, x1]
+ cr_loss = nn.functional.kl_div(
+ input=ctc_output,
+ target=exchanged_targets,
+ reduction="none",
+ log_target=True,
+ ) # (2 * N, T, C)
+ length_mask = make_pad_mask(encoder_out_lens).unsqueeze(-1)
+ cr_loss = cr_loss.masked_fill(length_mask, 0.0).sum()
+
+ return ctc_loss, cr_loss
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ y_asr: k2.RaggedTensor,
+ y_st: Optional[k2.RaggedTensor] = None,
+ srt_lang_ids: Optional[torch.Tensor] = None,
+ tgt_lang_ids: Optional[torch.Tensor] = None,
+ prune_range_asr: int = 5,
+ prune_range_st: int = 10,
+ am_scale: float = 0.0,
+ lm_scale: float = 0.0,
+ use_cr_ctc: bool = True,
+ use_spec_aug: bool = False,
+ spec_augment: Optional[SpecAugment] = None,
+ supervision_segments: Optional[torch.Tensor] = None,
+ time_warp_factor: Optional[int] = 80,
+ enable_st: bool = True,
+ ):
+
+ assert x.ndim == 3, x.shape
+ assert x_lens.ndim == 1, x_lens.shape
+ assert y_asr.num_axes == 2, y_asr.num_axes
+ if enable_st:
+ assert y_st is not None
+ assert y_st.num_axes == 2, y_st.num_axes
+ assert x.size(0) == x_lens.size(0) == y_asr.dim0 == y_st.dim0, (
+ x.shape,
+ x_lens.shape,
+ y_asr.dim0,
+ y_st.dim0,
+ )
+ else:
+ assert x.size(0) == x_lens.size(0) == y_asr.dim0
+ device = x.device
+
+ if use_cr_ctc:
+ assert self.use_ctc_asr
+ if use_spec_aug:
+ assert spec_augment is not None and spec_augment.time_warp_factor < 1
+ # Apply time warping before input duplicating
+ assert supervision_segments is not None
+ x = time_warp(
+ x,
+ time_warp_factor=time_warp_factor,
+ supervision_segments=supervision_segments,
+ )
+ # Independently apply frequency masking and time masking to the two copies
+ x = spec_augment(x.repeat(2, 1, 1))
+ else:
+ x = x.repeat(2, 1, 1)
+
+ x_lens = x_lens.repeat(2)
+ y_asr = k2.ragged.cat([y_asr, y_asr], axis=0)
+
+ if enable_st and (y_st is not None):
+ y_st = k2.ragged.cat([y_st, y_st], axis=0)
+
+ if srt_lang_ids is not None:
+ srt_lang_ids = srt_lang_ids.repeat(2)
+ srt_lang_ids = srt_lang_ids.to(x.device).view(-1)
+ if tgt_lang_ids is not None and enable_st:
+ tgt_lang_ids = tgt_lang_ids.repeat(2)
+ tgt_lang_ids = tgt_lang_ids.to(x.device).view(-1)
+
+ asr_ntc, asr_lens, st_ntc, st_lens, moe_ent_loss = self.forward_encoder(
+ x,
+ x_lens,
+ srt_lang_ids=srt_lang_ids,
+ tgt_lang_ids=tgt_lang_ids if enable_st else None,
+ enable_st=enable_st,
+ )
+
+ # ASR prepare targets
+ row_splits_asr = y_asr.shape.row_splits(1)
+ y_lens_asr = row_splits_asr[1:] - row_splits_asr[:-1]
+
+ # AST prepare targets
+ if enable_st and (y_st is not None):
+ row_splits_st = y_st.shape.row_splits(1)
+ y_lens_st = row_splits_st[1:] - row_splits_st[:-1]
+ else:
+ y_lens_st = None
+
+ # ---------- RNNT losses ----------
+ if self.use_transducer:
+ simple_loss_asr, pruned_loss_asr = self.forward_transducer(
+ encoder_out=asr_ntc,
+ encoder_out_lens=asr_lens,
+ y=y_asr.to(device),
+ y_lens=y_lens_asr,
+ decoder=self.decoder_asr,
+ joiner=self.joiner_asr,
+ simple_lm_proj=self.simple_lm_proj_asr,
+ simple_am_proj=self.simple_am_proj_asr,
+ prune_range=prune_range_asr,
+ am_scale=am_scale,
+ lm_scale=lm_scale,
+ )
+
+ if enable_st and (y_st is not None) and (y_lens_st is not None):
+ simple_loss_st, pruned_loss_st = self.forward_transducer(
+ encoder_out=st_ntc,
+ encoder_out_lens=st_lens,
+ y=y_st.to(device),
+ y_lens=y_lens_st,
+ decoder=self.decoder_st,
+ joiner=self.joiner_st,
+ simple_lm_proj=self.simple_lm_proj_st,
+ simple_am_proj=self.simple_am_proj_st,
+ prune_range=prune_range_st,
+ am_scale=am_scale,
+ lm_scale=lm_scale,
+ )
+ else:
+ simple_loss_st = pruned_loss_st = torch.empty(0)
+
+ if use_cr_ctc:
+ simple_loss_asr = simple_loss_asr * 0.5
+ pruned_loss_asr = pruned_loss_asr * 0.5
+ if simple_loss_st.numel() > 0:
+ simple_loss_st = simple_loss_st * 0.5
+ pruned_loss_st = pruned_loss_st * 0.5
+ else:
+ simple_loss_asr = pruned_loss_asr = torch.empty(0)
+ simple_loss_st = pruned_loss_st = torch.empty(0)
+
+ if self.use_ctc_asr:
+ # Compute CTC loss
+ targets_asr = y_asr.values
+
+ if not use_cr_ctc:
+ ctc_loss_asr = self.forward_ctc(
+ encoder_out=asr_ntc,
+ encoder_out_lens=asr_lens,
+ targets=targets_asr,
+ target_lengths=y_lens_asr,
+ ctc_output=self.ctc_asr,
+ )
+ cr_loss_asr = torch.empty(0)
+ else:
+ ctc_loss_asr, cr_loss_asr = self.forward_cr_ctc(
+ encoder_out=asr_ntc,
+ encoder_out_lens=asr_lens,
+ targets=targets_asr,
+ target_lengths=y_lens_asr,
+ ctc_output=self.ctc_asr,
+ )
+ ctc_loss_asr = ctc_loss_asr * 0.5
+ cr_loss_asr = cr_loss_asr * 0.5
+ else:
+ ctc_loss_asr = cr_loss_asr = torch.empty(0)
+
+ if (
+ self.use_ctc_st
+ and enable_st
+ and (y_st is not None)
+ and (y_lens_st is not None)
+ ):
+ targets_st = y_st.values
+ if not use_cr_ctc:
+ ctc_loss_st = self.forward_ctc(
+ encoder_out=st_ntc,
+ encoder_out_lens=st_lens,
+ targets=targets_st,
+ target_lengths=y_lens_st,
+ ctc_output=self.ctc_st,
+ )
+ cr_loss_st = torch.empty(0)
+ else:
+ ctc_loss_st, cr_loss_st = self.forward_cr_ctc(
+ encoder_out=st_ntc,
+ encoder_out_lens=st_lens,
+ targets=targets_st,
+ target_lengths=y_lens_st,
+ ctc_output=self.ctc_st,
+ )
+ ctc_loss_st = ctc_loss_st * 0.5
+ cr_loss_st = cr_loss_st * 0.5
+ else:
+ ctc_loss_st = cr_loss_st = torch.empty(0)
+
+ if self.use_attention_decoder:
+ attention_decoder_loss_asr = self.attention_decoder_asr.calc_att_loss(
+ encoder_out=asr_ntc,
+ encoder_out_lens=asr_lens,
+ ys=y_asr.to(device),
+ ys_lens=y_lens_asr.to(device),
+ )
+ if use_cr_ctc:
+ attention_decoder_loss_asr = attention_decoder_loss_asr * 0.5
+
+ if enable_st and (y_st is not None) and (y_lens_st is not None):
+ attention_decoder_loss_st = self.attention_decoder_st.calc_att_loss(
+ encoder_out=st_ntc,
+ encoder_out_lens=st_lens,
+ ys=y_st.to(device),
+ ys_lens=y_lens_st.to(device),
+ )
+ if use_cr_ctc:
+ attention_decoder_loss_st = attention_decoder_loss_st * 0.5
+ else:
+ attention_decoder_loss_st = torch.empty(0)
+ else:
+ attention_decoder_loss_asr = torch.empty(0)
+ attention_decoder_loss_st = torch.empty(0)
+
+ return (
+ simple_loss_asr,
+ simple_loss_st,
+ pruned_loss_asr,
+ pruned_loss_st,
+ ctc_loss_asr,
+ ctc_loss_st,
+ attention_decoder_loss_asr,
+ attention_decoder_loss_st,
+ cr_loss_asr,
+ cr_loss_st,
+ moe_ent_loss,
+ )
diff --git a/egs/europarl_st/SRT/lcma_srt/moe_adapter.py b/egs/europarl_st/SRT/lcma_srt/moe_adapter.py
new file mode 100644
index 0000000000..50adfc4f60
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/moe_adapter.py
@@ -0,0 +1,207 @@
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+class MoEAdapterDense(nn.Module):
+ def __init__(
+ self,
+ d_model: int,
+ num_experts: int = 4,
+ hidden_mult: float = 1.3,
+ num_tasks: int = 2,
+ num_langs: int = 0,
+ *,
+ num_src_langs: int = 0,
+ num_tgt_langs: int = 0,
+ dropout: float = 0.1,
+ router_mid_ratio: float = 0.5,
+ entropy_reg: float = 0.0,
+ temperature: float = 1.0,
+ init_identity: bool = True,
+ topk_infer: Optional[int] = None,
+ in_dropout: float = 0.0,
+ ):
+ super().__init__()
+ assert temperature > 0.0, "temperature must be > 0"
+ assert num_experts >= 1, "num_experts must be >= 1"
+
+ self.d_model = d_model
+ self.num_experts = num_experts
+ self.entropy_reg = float(entropy_reg)
+ self.temperature = float(temperature)
+ self.topk_infer = topk_infer
+
+ # --- Experts: MLP(d -> d_hidden -> d) + residual outside ---
+ d_hidden = max(1, int(round(d_model * hidden_mult)))
+ self.experts = nn.ModuleList(
+ [
+ nn.Sequential(
+ nn.Linear(d_model, d_hidden),
+ nn.ReLU(inplace=False),
+ nn.Dropout(dropout),
+ nn.Linear(d_hidden, d_model),
+ )
+ for _ in range(num_experts)
+ ]
+ )
+ if init_identity:
+ for e in self.experts:
+ nn.init.zeros_(e[-1].weight)
+ if e[-1].bias is not None:
+ nn.init.zeros_(e[-1].bias)
+
+ # --- Router and conditional embeddings ---
+ router_hidden = max(16, int(round(d_model * router_mid_ratio)))
+ self.task_embed = nn.Embedding(num_tasks, d_model) if num_tasks > 0 else None
+ self.lang_embed = nn.Embedding(num_langs, d_model) if num_langs > 0 else None
+ if num_src_langs > 0 or num_tgt_langs > 0:
+ assert num_langs == 0, "num_langs must be 0 when using src+tgt embeds"
+ self.lang_embed_src = (
+ nn.Embedding(num_src_langs, d_model) if num_src_langs > 0 else None
+ )
+ self.lang_embed_tgt = (
+ nn.Embedding(num_tgt_langs, d_model) if num_tgt_langs > 0 else None
+ )
+
+ router_in_dim = (
+ d_model
+ + (d_model if self.task_embed is not None else 0)
+ + (d_model if self.lang_embed is not None else 0)
+ + (d_model if self.lang_embed_src is not None else 0)
+ + (d_model if self.lang_embed_tgt is not None else 0)
+ )
+
+ self.router = nn.Sequential(
+ nn.Linear(router_in_dim, router_hidden),
+ nn.ReLU(inplace=False),
+ nn.Linear(router_hidden, num_experts),
+ )
+
+ self.in_drop = nn.Dropout(in_dropout)
+ self.last_router_weights: Optional[Tensor] = None
+
+ # ---- Helpers -----------------------------------------------------------------
+ def _apply_topk(self, w: Tensor) -> Tensor:
+ if (self.topk_infer is None) or self.training:
+ return w
+ k = min(self.topk_infer, w.size(-1))
+ topv, topi = torch.topk(w, k, dim=-1)
+ mask = torch.zeros_like(w).scatter_(-1, topi, topv)
+ return mask / (mask.sum(-1, keepdim=True) + 1e-9)
+
+ # ---- Forward -----------------------------------------------------------------
+ def forward(
+ self,
+ x_tbc: Tensor,
+ task_id: Optional[int] = None,
+ lang_ids: Optional[Tensor] = None,
+ src_lang_ids: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Tensor]:
+ T, B, C = x_tbc.shape
+ assert (
+ C == self.d_model
+ ), f"Mismatched channel: got {C}, expected {self.d_model}"
+ dtype = x_tbc.dtype
+ device = x_tbc.device
+
+ conds = [x_tbc]
+
+ if self.task_embed is not None:
+ assert task_id is not None, "task_id must be provided when num_tasks > 0"
+ t_id = torch.tensor(int(task_id), device=device, dtype=torch.long)
+ tvec = self.task_embed(t_id).view(1, 1, C).expand(T, B, C).to(dtype)
+ conds.append(tvec)
+
+ if self.lang_embed is not None:
+ if lang_ids is None:
+ raise ValueError("lang_ids must be provided when num_langs > 0")
+ if lang_ids.dtype != torch.long:
+ lang_ids = lang_ids.long()
+ lvec = (
+ self.lang_embed(lang_ids.to(device))
+ .view(1, B, C)
+ .expand(T, B, C)
+ .to(dtype)
+ )
+ conds.append(lvec)
+ if self.lang_embed_src is not None or self.lang_embed_tgt is not None:
+ if self.lang_embed_src is not None:
+ if src_lang_ids is None:
+ raise ValueError(
+ "src_lang_ids must be provided when src+tgt embeds are enabled"
+ )
+ sids = src_lang_ids.to(device)
+ if sids.dtype != torch.long:
+ sids = sids.long()
+ svec = self.lang_embed_src(sids).view(1, B, C).expand(T, B, C).to(dtype)
+ conds.append(svec)
+ if self.lang_embed_tgt is not None:
+ if lang_ids is None:
+ raise ValueError(
+ "lang_ids (tgt) must be provided when src+tgt embeds are enabled"
+ )
+ tids = lang_ids.to(device)
+ if tids.dtype != torch.long:
+ tids = tids.long()
+ tvec = self.lang_embed_tgt(tids).view(1, B, C).expand(T, B, C).to(dtype)
+ conds.append(tvec)
+
+ r_in = torch.cat(conds, dim=-1) # [T, B, Cin]
+ logits = self.router(r_in.float()) / self.temperature
+ w = torch.softmax(logits, dim=-1).to(dtype) # [T, B, E]
+ w = self._apply_topk(w)
+ self.last_router_weights = w
+
+ x_in = self.in_drop(x_tbc)
+ E = torch.stack([e(x_in) for e in self.experts], dim=-1)
+ # y = sum_e w_e * expert_e(x)
+ y = torch.einsum("tbce,tbe->tbc", E, w)
+ return x_tbc + y, w
+
+ # ---- Regularization -----------------------------------------------------------
+ def router_entropy_loss(
+ self,
+ w: Optional[Tensor] = None,
+ pad_mask_tb: Optional[Tensor] = None,
+ ) -> Tensor:
+ if self.entropy_reg <= 0.0:
+ p = next(self.parameters())
+ return torch.zeros((), device=p.device, dtype=p.dtype)
+
+ if w is None:
+ w = self.last_router_weights
+ if w is None:
+ p = next(self.parameters())
+ return torch.zeros((), device=p.device, dtype=p.dtype)
+
+ w_safe = w.clamp_min(1e-9)
+ ent_tb = -(w * w_safe.log()).sum(dim=-1) # [T,B]
+
+ if pad_mask_tb is not None:
+ if pad_mask_tb.dtype != torch.bool:
+ pad_mask_tb = pad_mask_tb.bool()
+ if pad_mask_tb.device != ent_tb.device:
+ pad_mask_tb = pad_mask_tb.to(ent_tb.device)
+ valid = (~pad_mask_tb).float()
+ denom = valid.sum().clamp_min(1.0)
+ ent = (ent_tb * valid).sum() / denom
+ else:
+ ent = ent_tb.mean()
+ return -self.entropy_reg * ent
diff --git a/egs/europarl_st/SRT/lcma_srt/optim.py b/egs/europarl_st/SRT/lcma_srt/optim.py
new file mode 100644
index 0000000000..8a17646513
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/optim.py
@@ -0,0 +1,1237 @@
+# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey)
+#
+# See ../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import logging
+import random
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from lhotse.utils import fix_random_seed
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class BatchedOptimizer(Optimizer):
+ """
+ This class adds to class Optimizer the capability to optimize parameters in batches:
+ it will stack the parameters and their grads for you so the optimizer can work
+ on tensors with an extra leading dimension. This is intended for speed with GPUs,
+ as it reduces the number of kernels launched in the optimizer.
+
+ Args:
+ params:
+ """
+
+ def __init__(self, params, defaults):
+ super(BatchedOptimizer, self).__init__(params, defaults)
+
+ @contextlib.contextmanager
+ def batched_params(self, param_group, group_params_names):
+ """
+ This function returns (technically, yields) a list of
+ of tuples (p, state), where
+ p is a `fake` parameter that is stacked (over axis 0) from real parameters
+ that share the same shape, and its gradient is also stacked;
+ `state` is the state corresponding to this batch of parameters
+ (it will be physically located in the "state" for one of the real
+ parameters, the last one that has any particular shape and dtype).
+
+ This function is decorated as a context manager so that it can
+ write parameters back to their "real" locations.
+
+ The idea is, instead of doing:
+
+ for p in group["params"]:
+ state = self.state[p]
+ ...
+
+ you can do:
+
+ with self.batched_params(group["params"]) as batches:
+ for p, state, p_names in batches:
+ ...
+
+
+ Args:
+ group: a parameter group, which is a list of parameters; should be
+ one of self.param_groups.
+ group_params_names: name for each parameter in group,
+ which is List[str].
+ """
+ batches = defaultdict(
+ list
+ ) # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
+ batches_names = defaultdict(
+ list
+ ) # `batches` maps from tuple (dtype_as_str,*shape) to list of str
+
+ assert len(param_group) == len(group_params_names)
+ for p, named_p in zip(param_group, group_params_names):
+ key = (str(p.dtype), *p.shape)
+ batches[key].append(p)
+ batches_names[key].append(named_p)
+
+ batches_names_keys = list(batches_names.keys())
+ sorted_idx = sorted(
+ range(len(batches_names)), key=lambda i: batches_names_keys[i]
+ )
+ batches_names = [batches_names[batches_names_keys[idx]] for idx in sorted_idx]
+ batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
+
+ stacked_params_dict = dict()
+
+ # turn batches into a list, in deterministic order.
+ # tuples will contain tuples of (stacked_param, state, stacked_params_names),
+ # one for each batch in `batches`.
+ tuples = []
+
+ for batch, batch_names in zip(batches, batches_names):
+ p = batch[0]
+ # we arbitrarily store the state in the
+ # state corresponding to the 1st parameter in the
+ # group. class Optimizer will take care of saving/loading state.
+ state = self.state[p]
+ p_stacked = torch.stack(batch)
+ grad = torch.stack(
+ [torch.zeros_like(p) if p.grad is None else p.grad for p in batch]
+ )
+ p_stacked.grad = grad
+ stacked_params_dict[key] = p_stacked
+ tuples.append((p_stacked, state, batch_names))
+
+ yield tuples # <-- calling code will do the actual optimization here!
+
+ for ((stacked_params, _state, _names), batch) in zip(tuples, batches):
+ for i, p in enumerate(batch): # batch is list of Parameter
+ p.copy_(stacked_params[i])
+
+
+def basic_step(group, p, state, grad):
+ # computes basic Adam update using beta2 (dividing by gradient stddev) only. no momentum yet.
+ lr = group["lr"]
+ if p.numel() == p.shape[0]:
+ lr = lr * group["scalar_lr_scale"]
+ beta2 = group["betas"][1]
+ eps = group["eps"]
+ # p shape: (batch_size,) or (batch_size, 1, [1,..])
+ try:
+ exp_avg_sq = state[
+ "exp_avg_sq"
+ ] # shape: (batch_size,) or (batch_size, 1, [1,..])
+ except KeyError:
+ exp_avg_sq = torch.zeros(*p.shape, device=p.device, dtype=torch.float)
+ state["exp_avg_sq"] = exp_avg_sq
+
+ exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+ # bias_correction2 is like in Adam.
+ # slower update at the start will help stability anyway.
+ bias_correction2 = 1 - beta2 ** (state["step"] + 1)
+ if bias_correction2 < 0.99:
+ # note: not in-place.
+ exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
+ denom = exp_avg_sq.sqrt().add_(eps)
+
+ return -lr * grad / denom
+
+
+def scaling_step(group, p, state, grad):
+ delta = basic_step(group, p, state, grad)
+ if p.numel() == p.shape[0]:
+ return delta # there is no scaling for scalar parameters. (p.shape[0] is the batch of parameters.)
+
+ step = state["step"]
+ size_update_period = group["size_update_period"]
+
+ try:
+ param_rms = state["param_rms"]
+ scale_grads = state["scale_grads"]
+ scale_exp_avg_sq = state["scale_exp_avg_sq"]
+ except KeyError:
+ # we know p.ndim > 1 because we'd have returned above if not, so don't worry
+ # about the speial case of dim=[] that pytorch treats inconsistently.
+ param_rms = (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()
+ param_rms = param_rms.to(torch.float)
+ scale_exp_avg_sq = torch.zeros_like(param_rms)
+ scale_grads = torch.zeros(
+ size_update_period, *param_rms.shape, dtype=torch.float, device=p.device
+ )
+ state["param_rms"] = param_rms
+ state["scale_grads"] = scale_grads
+ state["scale_exp_avg_sq"] = scale_exp_avg_sq
+
+ # on every step, update the gradient w.r.t. the scale of the parameter, we
+ # store these as a batch and periodically update the size (for speed only, to
+ # avoid too many operations).
+ scale_grads[step % size_update_period] = (p * grad).sum(
+ dim=list(range(1, p.ndim)), keepdim=True
+ )
+
+ # periodically recompute the value of param_rms.
+ if step % size_update_period == size_update_period - 1:
+ param_rms.copy_((p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt())
+
+ param_min_rms = group["param_min_rms"]
+
+ # scale the step size by param_rms. This is the most important "scaling" part of
+ # ScaledAdam
+ delta *= param_rms.clamp(min=param_min_rms)
+
+ if step % size_update_period == size_update_period - 1 and step > 0:
+ # This block updates the size of parameter by adding a step ("delta") value in
+ # the direction of either shrinking or growing it.
+ beta2 = group["betas"][1]
+ size_lr = group["lr"] * group["scalar_lr_scale"]
+ param_max_rms = group["param_max_rms"]
+ eps = group["eps"]
+ batch_size = p.shape[0]
+ # correct beta2 for the size update period: we will have
+ # faster decay at this level.
+ beta2_corr = beta2**size_update_period
+ scale_exp_avg_sq.mul_(beta2_corr).add_(
+ (scale_grads**2).mean(dim=0), # mean over dim `size_update_period`
+ alpha=1 - beta2_corr,
+ ) # shape is (batch_size, 1, 1, ...)
+
+ # The 1st time we reach here is when size_step == 1.
+ size_step = (step + 1) // size_update_period
+ bias_correction2 = 1 - beta2_corr**size_step
+
+ denom = scale_exp_avg_sq.sqrt() + eps
+
+ scale_step = (
+ -size_lr * (bias_correction2**0.5) * scale_grads.sum(dim=0) / denom
+ )
+
+ is_too_small = param_rms < param_min_rms
+
+ # when the param gets too small, just don't shrink it any further.
+ scale_step.masked_fill_(is_too_small, 0.0)
+
+ # The following may help prevent instability: don't allow the scale step to be too large in
+ # either direction.
+ scale_step.clamp_(min=-0.1, max=0.1)
+
+ # and ensure the parameter rms after update never exceeds param_max_rms.
+ # We have to look at the trained model for parameters at or around the
+ # param_max_rms, because sometimes they can indicate a problem with the
+ # topology or settings.
+ scale_step = torch.minimum(scale_step, (param_max_rms - param_rms) / param_rms)
+
+ delta.add_(p * scale_step)
+
+ return delta
+
+
+def momentum_step(group, p, state, grad):
+ delta = scaling_step(group, p, state, grad)
+ beta1 = group["betas"][0]
+ try:
+ stored_delta = state["delta"]
+ except KeyError:
+ stored_delta = torch.zeros(*p.shape, device=p.device, dtype=torch.float)
+ state["delta"] = stored_delta
+ stored_delta.mul_(beta1)
+ stored_delta.add_(delta, alpha=(1 - beta1))
+ # we don't bother doing the "bias correction" part of Adam for beta1 because this is just
+ # an edge effect that affects the first 10 or so batches; and the effect of not doing it
+ # is just to do a slower update for the first few batches, which will help stability.
+ return stored_delta
+
+
+class ScaledAdam(BatchedOptimizer):
+ """
+ Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
+ proportional to the norm of that parameter; and also learn the scale of the parameter,
+ in log space, subject to upper and lower limits (as if we had factored each parameter as
+ param = underlying_param * log_scale.exp())
+
+
+ Args:
+ params: The parameters or param_groups to optimize (like other Optimizer subclasses)
+ Unlike common optimizers, which accept model.parameters() or groups of parameters(),
+ this optimizer could accept model.named_parameters() or groups of named_parameters().
+ See comments of function _get_names_of_parameters for its 4 possible cases.
+ lr: The learning rate. We will typically use a learning rate schedule that starts
+ at 0.03 and decreases over time, i.e. much higher than other common
+ optimizers.
+ clipping_scale: (e.g. 2.0)
+ A scale for gradient-clipping: if specified, the normalized gradients
+ over the whole model will be clipped to have 2-norm equal to
+ `clipping_scale` times the median 2-norm over the most recent period
+ of `clipping_update_period` minibatches. By "normalized gradients",
+ we mean after multiplying by the rms parameter value for this tensor
+ [for non-scalars]; this is appropriate because our update is scaled
+ by this quantity.
+ betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.
+ Must satisfy 0 < beta <= beta2 < 1.
+ scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
+ scale of each parameter tensor and scalar parameters of the mode..
+ If each parameter were decomposed
+ as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale
+ would be a the scaling factor on the learning rate of p_scale.
+ eps: A general-purpose epsilon to prevent division by zero
+ param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
+ learning the scale on the parameters (we'll constrain the rms of each non-scalar
+ parameter tensor to be >= this value)
+ param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
+ learning the scale on the parameters (we'll constrain the rms of each non-scalar
+ parameter tensor to be <= this value)
+ scalar_max: Maximum absolute value for scalar parameters (applicable if your
+ model has any parameters with numel() == 1).
+ size_update_period: The periodicity, in steps, with which we update the size (scale)
+ of the parameter tensor. This is provided to save a little time
+ in the update.
+ clipping_update_period: if clipping_scale is specified, this is the period
+ """
+
+ def __init__(
+ self,
+ params,
+ lr=3e-02,
+ clipping_scale=None,
+ betas=(0.9, 0.98),
+ scalar_lr_scale=0.1,
+ eps=1.0e-08,
+ param_min_rms=1.0e-05,
+ param_max_rms=3.0,
+ scalar_max=10.0,
+ size_update_period=4,
+ clipping_update_period=100,
+ ):
+
+ defaults = dict(
+ lr=lr,
+ clipping_scale=clipping_scale,
+ betas=betas,
+ scalar_lr_scale=scalar_lr_scale,
+ eps=eps,
+ param_min_rms=param_min_rms,
+ param_max_rms=param_max_rms,
+ scalar_max=scalar_max,
+ size_update_period=size_update_period,
+ clipping_update_period=clipping_update_period,
+ )
+
+ # If params only contains parameters or group of parameters,
+ # i.e when parameter names are not given,
+ # this flag will be set to False in funciton _get_names_of_parameters.
+ self.show_dominant_parameters = True
+ param_groups, parameters_names = self._get_names_of_parameters(params)
+ super(ScaledAdam, self).__init__(param_groups, defaults)
+ assert len(self.param_groups) == len(parameters_names)
+ self.parameters_names = parameters_names
+
+ def _get_names_of_parameters(
+ self, params_or_named_params
+ ) -> Tuple[List[Dict], List[List[str]]]:
+ """
+ Args:
+ params_or_named_params: according to the way ScaledAdam is initialized in train.py,
+ this argument could be one of following 4 cases,
+ case 1, a generator of parameter, e.g.:
+ optimizer = ScaledAdam(model.parameters(), lr=params.base_lr, clipping_scale=3.0)
+
+ case 2, a list of parameter groups with different config, e.g.:
+ model_param_groups = [
+ {'params': model.encoder.parameters(), 'lr': 0.05},
+ {'params': model.decoder.parameters(), 'lr': 0.01},
+ {'params': model.joiner.parameters(), 'lr': 0.03},
+ ]
+ optimizer = ScaledAdam(model_param_groups, lr=params.base_lr, clipping_scale=3.0)
+
+ case 3, a generator of named_parameter, e.g.:
+ optimizer = ScaledAdam(model.named_parameters(), lr=params.base_lr, clipping_scale=3.0)
+
+ case 4, a list of named_parameter groups with different config, e.g.:
+ model_named_param_groups = [
+ {'named_params': model.encoder.named_parameters(), 'lr': 0.05},
+ {'named_params': model.decoder.named_parameters(), 'lr': 0.01},
+ {'named_params': model.joiner.named_parameters(), 'lr': 0.03},
+ ]
+ optimizer = ScaledAdam(model_named_param_groups, lr=params.base_lr, clipping_scale=3.0)
+
+ For case 1 and case 2, input params is used to initialize the underlying torch.optimizer.
+ For case 3 and case 4, firstly, names and params are extracted from input named_params,
+ then, these extracted params are used to initialize the underlying torch.optimizer,
+ and these extracted names are mainly used by function
+ `_show_gradient_dominating_parameter`
+
+ Returns:
+ Returns a tuple containing 2 elements:
+ - `param_groups` with type List[Dict], each Dict element is a parameter group.
+ An example of `param_groups` could be:
+ [
+ {'params': `one iterable of Parameter`, 'lr': 0.05},
+ {'params': `another iterable of Parameter`, 'lr': 0.08},
+ {'params': `a third iterable of Parameter`, 'lr': 0.1},
+ ]
+ - `param_gruops_names` with type List[List[str]],
+ each `List[str]` is for a group['params'] in param_groups,
+ and each `str` is the name of a parameter.
+ A dummy name "foo" is related to each parameter,
+ if input are params without names, i.e. case 1 or case 2.
+ """
+ # variable naming convention in this function:
+ # p is short for param.
+ # np is short for named_param.
+ # p_or_np is short for param_or_named_param.
+ # cur is short for current.
+ # group is a dict, e.g. {'params': iterable of parameter, 'lr': 0.05, other fields}.
+ # groups is a List[group]
+
+ iterable_or_groups = list(params_or_named_params)
+ if len(iterable_or_groups) == 0:
+ raise ValueError("optimizer got an empty parameter list")
+
+ # The first value of returned tuple. A list of dicts containing at
+ # least 'params' as a key.
+ param_groups = []
+
+ # The second value of returned tuple,
+ # a List[List[str]], each sub-List is for a group.
+ param_groups_names = []
+
+ if not isinstance(iterable_or_groups[0], dict):
+ # case 1 or case 3,
+ # the input is an iterable of parameter or named parameter.
+ param_iterable_cur_group = []
+ param_names_cur_group = []
+ for p_or_np in iterable_or_groups:
+ if isinstance(p_or_np, tuple):
+ # case 3
+ name, param = p_or_np
+ else:
+ # case 1
+ assert isinstance(p_or_np, torch.Tensor)
+ param = p_or_np
+ # Assign a dummy name as a placeholder
+ name = "foo"
+ self.show_dominant_parameters = False
+ param_iterable_cur_group.append(param)
+ param_names_cur_group.append(name)
+ param_groups.append({"params": param_iterable_cur_group})
+ param_groups_names.append(param_names_cur_group)
+ else:
+ # case 2 or case 4
+ # the input is groups of parameter or named parameter.
+ for cur_group in iterable_or_groups:
+ if "named_params" in cur_group:
+ name_list = [x[0] for x in cur_group["named_params"]]
+ p_list = [x[1] for x in cur_group["named_params"]]
+ del cur_group["named_params"]
+ cur_group["params"] = p_list
+ else:
+ assert "params" in cur_group
+ name_list = ["foo" for _ in cur_group["params"]]
+ param_groups.append(cur_group)
+ param_groups_names.append(name_list)
+
+ return param_groups, param_groups_names
+
+ def __setstate__(self, state):
+ super(ScaledAdam, self).__setstate__(state)
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ """Performs a single optimization step.
+
+ Arguments:
+ closure (callable, optional): A closure that reevaluates the model
+ and returns the loss.
+ """
+ loss = None
+ if closure is not None:
+ with torch.enable_grad():
+ loss = closure()
+
+ batch = True
+
+ for group, group_params_names in zip(self.param_groups, self.parameters_names):
+
+ with self.batched_params(group["params"], group_params_names) as batches:
+
+ # batches is list of pairs (stacked_param, state). stacked_param is like
+ # a regular parameter, and will have a .grad, but the 1st dim corresponds to
+ # a stacking dim, it is not a real dim.
+
+ if (
+ len(batches[0][1]) == 0
+ ): # if len(first state) == 0: not yet initialized
+ clipping_scale = 1
+ else:
+ clipping_scale = self._get_clipping_scale(group, batches)
+
+ for p, state, _ in batches:
+ # Perform optimization step.
+ # grad is not going to be None, we handled that when creating the batches.
+ grad = p.grad
+ if grad.is_sparse:
+ raise RuntimeError(
+ "ScaledAdam optimizer does not support sparse gradients"
+ )
+
+ try:
+ cur_step = state["step"]
+ except KeyError:
+ state["step"] = 0
+ cur_step = 0
+
+ grad = (
+ p.grad if clipping_scale == 1.0 else p.grad.mul_(clipping_scale)
+ )
+ p += momentum_step(group, p.detach(), state, grad)
+
+ if p.numel() == p.shape[0]: # scalar parameter
+ scalar_max = group["scalar_max"]
+ p.clamp_(min=-scalar_max, max=scalar_max)
+
+ state["step"] = cur_step + 1
+
+ return loss
+
+ def _get_clipping_scale(
+ self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]]
+ ) -> float:
+ """
+ Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients
+ by this amount before applying the rest of the update.
+
+ Args:
+ group: the parameter group, an item in self.param_groups
+ tuples: a list of tuples of (param, state, param_names)
+ where param is a batched set of parameters,
+ with a .grad (1st dim is batch dim)
+ and state is the state-dict where optimization parameters are kept.
+ param_names is a List[str] while each str is name for a parameter
+ in batched set of parameters "param".
+ """
+ assert len(tuples) >= 1
+ clipping_scale = group["clipping_scale"]
+ (first_p, first_state, _) = tuples[0]
+ step = first_state["step"]
+ if clipping_scale is None or step == 0:
+ # no clipping. return early on step == 0 because the other
+ # parameters' state won't have been initialized yet.
+ return 1.0
+ clipping_update_period = group["clipping_update_period"]
+ scalar_lr_scale = group["scalar_lr_scale"]
+
+ tot_sumsq = torch.tensor(0.0, device=first_p.device)
+ for (p, state, param_names) in tuples:
+ grad = p.grad
+ if grad.is_sparse:
+ raise RuntimeError(
+ "ScaledAdam optimizer does not support sparse gradients"
+ )
+ if p.numel() == p.shape[0]: # a batch of scalars
+ tot_sumsq += (grad**2).sum() * (
+ scalar_lr_scale**2
+ ) # sum() to change shape [1] to []
+ else:
+ tot_sumsq += ((grad * state["param_rms"]) ** 2).sum()
+
+ tot_norm = tot_sumsq.sqrt()
+ if "model_norms" not in first_state:
+ first_state["model_norms"] = torch.zeros(
+ clipping_update_period, device=p.device
+ )
+ first_state["model_norms"][step % clipping_update_period] = tot_norm
+
+ irregular_estimate_steps = [
+ i for i in [10, 20, 40] if i < clipping_update_period
+ ]
+ if step % clipping_update_period == 0 or step in irregular_estimate_steps:
+ # Print some stats.
+ # We don't reach here if step == 0 because we would have returned
+ # above.
+ sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
+ if step in irregular_estimate_steps:
+ sorted_norms = sorted_norms[-step:]
+ num_norms = sorted_norms.numel()
+ quartiles = []
+ for n in range(0, 5):
+ index = min(num_norms - 1, (num_norms // 4) * n)
+ quartiles.append(sorted_norms[index].item())
+
+ median = quartiles[2]
+ if median - median != 0:
+ raise RuntimeError("Too many grads were not finite")
+ threshold = clipping_scale * median
+ if step in irregular_estimate_steps:
+ # use larger thresholds on first few steps of estimating threshold,
+ # as norm may be changing rapidly.
+ threshold = threshold * 2.0
+ first_state["model_norm_threshold"] = threshold
+ percent_clipped = (
+ first_state["num_clipped"] * 100.0 / num_norms
+ if "num_clipped" in first_state
+ else 0.0
+ )
+ first_state["num_clipped"] = 0
+ quartiles = " ".join(["%.3e" % x for x in quartiles])
+ logging.warning(
+ f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
+ f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
+ )
+
+ try:
+ model_norm_threshold = first_state["model_norm_threshold"]
+ except KeyError:
+ return 1.0 # threshold has not yet been set.
+
+ ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
+ if ans != ans: # e.g. ans is nan
+ ans = 0.0
+ if ans < 1.0:
+ first_state["num_clipped"] += 1
+ if ans < 0.5:
+ logging.warning(
+ f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}"
+ )
+ if self.show_dominant_parameters:
+ assert p.shape[0] == len(param_names)
+ self._show_gradient_dominating_parameter(
+ tuples, tot_sumsq, group["scalar_lr_scale"]
+ )
+ self._show_param_with_unusual_grad(tuples)
+
+ if ans == 0.0:
+ for (p, state, param_names) in tuples:
+ p.grad.zero_() # get rid of infinity()
+
+ return ans
+
+ def _show_param_with_unusual_grad(
+ self,
+ tuples: List[Tuple[Tensor, dict, List[str]]],
+ ):
+ """
+ Print information about parameter which has the largest ratio of grad-on-this-batch
+ divided by normal grad size.
+ tuples: a list of tuples of (param, state, param_names)
+ where param is a batched set of parameters,
+ with a .grad (1st dim is batch dim)
+ and state is the state-dict where optimization parameters are kept.
+ param_names is a List[str] while each str is name for a parameter
+ in batched set of parameters "param".
+ """
+ largest_ratio = 0.0
+ largest_name = ""
+ # ratios_names is a list of 3-tuples: (grad_ratio, param_name, tensor)
+ ratios_names = []
+ for (p, state, batch_param_names) in tuples:
+ dims = list(range(1, p.ndim))
+
+ def mean(x):
+ # workaround for bad interface of torch's "mean" for when dims is the empty list.
+ if len(dims) > 0:
+ return x.mean(dim=dims)
+ else:
+ return x
+
+ grad_ratio = (
+ (mean(p.grad**2) / state["exp_avg_sq"].mean(dim=dims))
+ .sqrt()
+ .to("cpu")
+ )
+
+ ratios_names += zip(
+ grad_ratio.tolist(), batch_param_names, p.grad.unbind(dim=0)
+ )
+
+ ratios_names = sorted(ratios_names, reverse=True)
+ ratios_names = ratios_names[:10]
+ ratios_names = [
+ (ratio, name, largest_index(tensor))
+ for (ratio, name, tensor) in ratios_names
+ ]
+
+ logging.warning(
+ f"Parameters with most larger-than-usual grads, with ratios, are: {ratios_names}"
+ )
+
+ def _show_gradient_dominating_parameter(
+ self,
+ tuples: List[Tuple[Tensor, dict, List[str]]],
+ tot_sumsq: Tensor,
+ scalar_lr_scale: float,
+ ):
+ """
+ Show information of parameter which dominates tot_sumsq.
+
+ Args:
+ tuples: a list of tuples of (param, state, param_names)
+ where param is a batched set of parameters,
+ with a .grad (1st dim is batch dim)
+ and state is the state-dict where optimization parameters are kept.
+ param_names is a List[str] while each str is name for a parameter
+ in batched set of parameters "param".
+ tot_sumsq: sumsq of all parameters. Though it's could be calculated
+ from tuples, we still pass it to save some time.
+ """
+ all_sumsq_orig = {}
+ for (p, state, batch_param_names) in tuples:
+ # p is a stacked batch parameters.
+ batch_grad = p.grad
+ if p.numel() == p.shape[0]: # a batch of scalars
+ # Dummy values used by following `zip` statement.
+ batch_rms_orig = torch.full(
+ p.shape, scalar_lr_scale, device=batch_grad.device
+ )
+ else:
+ batch_rms_orig = state["param_rms"]
+ batch_sumsq_orig = (batch_grad * batch_rms_orig) ** 2
+ if batch_grad.ndim > 1:
+ # need to guard it with if-statement because sum() sums over
+ # all dims if dim == ().
+ batch_sumsq_orig = batch_sumsq_orig.sum(
+ dim=list(range(1, batch_grad.ndim))
+ )
+ for name, sumsq_orig, rms, grad in zip(
+ batch_param_names, batch_sumsq_orig, batch_rms_orig, batch_grad
+ ):
+
+ proportion_orig = sumsq_orig / tot_sumsq
+ all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
+
+ sorted_by_proportion = {
+ k: v
+ for k, v in sorted(
+ all_sumsq_orig.items(), key=lambda item: item[1][0], reverse=True
+ )
+ }
+ dominant_param_name = next(iter(sorted_by_proportion))
+ (
+ dominant_proportion,
+ dominant_sumsq,
+ dominant_rms,
+ dominant_grad,
+ ) = sorted_by_proportion[dominant_param_name]
+ logging.warning(
+ f"Parameter dominating tot_sumsq {dominant_param_name}"
+ f" with proportion {dominant_proportion:.2f},"
+ f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
+ f"={dominant_sumsq:.3e},"
+ f" grad_sumsq={(dominant_grad**2).sum():.3e},"
+ f" orig_rms_sq={(dominant_rms**2).item():.3e}"
+ )
+
+
+def largest_index(x: Tensor):
+ x = x.contiguous()
+ argmax = x.abs().argmax().item()
+ return [(argmax // x.stride(i)) % x.size(i) for i in range(x.ndim)]
+
+
+class LRScheduler(object):
+ """
+ Base-class for learning rate schedulers where the learning-rate depends on both the
+ batch and the epoch.
+ """
+
+ def __init__(self, optimizer: Optimizer, verbose: bool = False):
+ # Attach optimizer
+ if not isinstance(optimizer, Optimizer):
+ raise TypeError("{} is not an Optimizer".format(type(optimizer).__name__))
+ self.optimizer = optimizer
+ self.verbose = verbose
+
+ for group in optimizer.param_groups:
+ group.setdefault("base_lr", group["lr"])
+
+ self.base_lrs = [group["base_lr"] for group in optimizer.param_groups]
+
+ self.epoch = 0
+ self.batch = 0
+
+ def state_dict(self):
+ """Returns the state of the scheduler as a :class:`dict`.
+
+ It contains an entry for every variable in self.__dict__ which
+ is not the optimizer.
+ """
+ return {
+ # the user might try to override the base_lr, so don't include this in the state.
+ # previously they were included.
+ # "base_lrs": self.base_lrs,
+ "epoch": self.epoch,
+ "batch": self.batch,
+ }
+
+ def load_state_dict(self, state_dict):
+ """Loads the schedulers state.
+
+ Args:
+ state_dict (dict): scheduler state. Should be an object returned
+ from a call to :meth:`state_dict`.
+ """
+ # the things with base_lrs are a work-around for a previous problem
+ # where base_lrs were written with the state dict.
+ base_lrs = self.base_lrs
+ self.__dict__.update(state_dict)
+ self.base_lrs = base_lrs
+
+ def get_last_lr(self) -> List[float]:
+ """Return last computed learning rate by current scheduler. Will be a list of float."""
+ return self._last_lr
+
+ def get_lr(self):
+ # Compute list of learning rates from self.epoch and self.batch and
+ # self.base_lrs; this must be overloaded by the user.
+ # e.g. return [some_formula(self.batch, self.epoch, base_lr) for base_lr in self.base_lrs ]
+ raise NotImplementedError
+
+ def step_batch(self, batch: Optional[int] = None) -> None:
+ # Step the batch index, or just set it. If `batch` is specified, it
+ # must be the batch index from the start of training, i.e. summed over
+ # all epochs.
+ # You can call this in any order; if you don't provide 'batch', it should
+ # of course be called once per batch.
+ if batch is not None:
+ self.batch = batch
+ else:
+ self.batch = self.batch + 1
+ self._set_lrs()
+
+ def step_epoch(self, epoch: Optional[int] = None):
+ # Step the epoch index, or just set it. If you provide the 'epoch' arg,
+ # you should call this at the start of the epoch; if you don't provide the 'epoch'
+ # arg, you should call it at the end of the epoch.
+ if epoch is not None:
+ self.epoch = epoch
+ else:
+ self.epoch = self.epoch + 1
+ self._set_lrs()
+
+ def _set_lrs(self):
+ values = self.get_lr()
+ assert len(values) == len(self.optimizer.param_groups)
+
+ for i, data in enumerate(zip(self.optimizer.param_groups, values)):
+ param_group, lr = data
+ param_group["lr"] = lr
+ self.print_lr(self.verbose, i, lr)
+ self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
+
+ def print_lr(self, is_verbose, group, lr):
+ """Display the current learning rate."""
+ if is_verbose:
+ logging.warning(
+ f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
+ f" of group {group} to {lr:.4e}."
+ )
+
+
+class Eden(LRScheduler):
+ """
+ Eden scheduler.
+ The basic formula (before warmup) is:
+ lr = base_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
+ (((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25)) * warmup
+ where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches
+ and then stays constant at 1.
+
+ If you don't have the concept of epochs, or one epoch takes a very long time,
+ you can replace the notion of 'epoch' with some measure of the amount of data
+ processed, e.g. hours of data or frames of data, with 'lr_epochs' being set to
+ some measure representing "quite a lot of data": say, one fifth or one third
+ of an entire training run, but it doesn't matter much. You could also use
+ Eden2 which has only the notion of batches.
+
+ We suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam
+
+ Args:
+ optimizer: the optimizer to change the learning rates on
+ lr_batches: the number of batches after which we start significantly
+ decreasing the learning rate, suggest 5000.
+ lr_epochs: the number of epochs after which we start significantly
+ decreasing the learning rate, suggest 6 if you plan to do e.g.
+ 20 to 40 epochs, but may need smaller number if dataset is huge
+ and you will do few epochs.
+ """
+
+ def __init__(
+ self,
+ optimizer: Optimizer,
+ lr_batches: Union[int, float],
+ lr_epochs: Union[int, float],
+ warmup_batches: Union[int, float] = 500.0,
+ warmup_start: float = 0.5,
+ verbose: bool = False,
+ ):
+ super(Eden, self).__init__(optimizer, verbose)
+ self.lr_batches = lr_batches
+ self.lr_epochs = lr_epochs
+ self.warmup_batches = warmup_batches
+
+ assert 0.0 <= warmup_start <= 1.0, warmup_start
+ self.warmup_start = warmup_start
+
+ def get_lr(self):
+ factor = (
+ (self.batch**2 + self.lr_batches**2) / self.lr_batches**2
+ ) ** -0.25 * (
+ ((self.epoch**2 + self.lr_epochs**2) / self.lr_epochs**2) ** -0.25
+ )
+ warmup_factor = (
+ 1.0
+ if self.batch >= self.warmup_batches
+ else self.warmup_start
+ + (1.0 - self.warmup_start) * (self.batch / self.warmup_batches)
+ # else 0.5 + 0.5 * (self.batch / self.warmup_batches)
+ )
+
+ return [x * factor * warmup_factor for x in self.base_lrs]
+
+
+class Eden2(LRScheduler):
+ """
+ Eden2 scheduler, simpler than Eden because it does not use the notion of epoch,
+ only batches.
+
+ The basic formula (before warmup) is:
+ lr = base_lr * ((batch**2 + lr_batches**2) / lr_batches**2) ** -0.5) * warmup
+
+ where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches
+ and then stays constant at 1.
+
+
+ E.g. suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam
+
+ Args:
+ optimizer: the optimizer to change the learning rates on
+ lr_batches: the number of batches after which we start significantly
+ decreasing the learning rate, suggest 5000.
+ """
+
+ def __init__(
+ self,
+ optimizer: Optimizer,
+ lr_batches: Union[int, float],
+ warmup_batches: Union[int, float] = 500.0,
+ warmup_start: float = 0.5,
+ verbose: bool = False,
+ ):
+ super().__init__(optimizer, verbose)
+ self.lr_batches = lr_batches
+ self.warmup_batches = warmup_batches
+
+ assert 0.0 <= warmup_start <= 1.0, warmup_start
+ self.warmup_start = warmup_start
+
+ def get_lr(self):
+ factor = (
+ (self.batch**2 + self.lr_batches**2) / self.lr_batches**2
+ ) ** -0.5
+ warmup_factor = (
+ 1.0
+ if self.batch >= self.warmup_batches
+ else self.warmup_start
+ + (1.0 - self.warmup_start) * (self.batch / self.warmup_batches)
+ # else 0.5 + 0.5 * (self.batch / self.warmup_batches)
+ )
+
+ return [x * factor * warmup_factor for x in self.base_lrs]
+
+
+def _test_eden():
+ m = torch.nn.Linear(100, 100)
+ optim = ScaledAdam(m.parameters(), lr=0.03)
+
+ scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True)
+
+ for epoch in range(10):
+ scheduler.step_epoch(epoch) # sets epoch to `epoch`
+
+ for step in range(20):
+ x = torch.randn(200, 100).detach()
+ x.requires_grad = True
+ y = m(x)
+ dy = torch.randn(200, 100).detach()
+ f = (y * dy).sum()
+ f.backward()
+
+ optim.step()
+ scheduler.step_batch()
+ optim.zero_grad()
+
+ logging.info(f"last lr = {scheduler.get_last_lr()}")
+ logging.info(f"state dict = {scheduler.state_dict()}")
+
+
+# This is included mostly as a baseline for ScaledAdam.
+class Eve(Optimizer):
+ """
+ Implements Eve algorithm. This is a modified version of AdamW with a special
+ way of setting the weight-decay / shrinkage-factor, which is designed to make the
+ rms of the parameters approach a particular target_rms (default: 0.1). This is
+ for use with networks with 'scaled' versions of modules (see scaling.py), which
+ will be close to invariant to the absolute scale on the parameter matrix.
+
+ The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+ The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+ Eve is unpublished so far.
+
+ Arguments:
+ params (iterable): iterable of parameters to optimize or dicts defining
+ parameter groups
+ lr (float, optional): learning rate (default: 1e-3)
+ betas (Tuple[float, float], optional): coefficients used for computing
+ running averages of gradient and its square (default: (0.9, 0.999))
+ eps (float, optional): term added to the denominator to improve
+ numerical stability (default: 1e-8)
+ weight_decay (float, optional): weight decay coefficient (default: 3e-4;
+ this value means that the weight would decay significantly after
+ about 3k minibatches. Is not multiplied by learning rate, but
+ is conditional on RMS-value of parameter being > target_rms.
+ target_rms (float, optional): target root-mean-square value of
+ parameters, if they fall below this we will stop applying weight decay.
+
+
+ .. _Adam: A Method for Stochastic Optimization:
+ https://arxiv.org/abs/1412.6980
+ .. _Decoupled Weight Decay Regularization:
+ https://arxiv.org/abs/1711.05101
+ .. _On the Convergence of Adam and Beyond:
+ https://openreview.net/forum?id=ryQu7f-RZ
+ """
+
+ def __init__(
+ self,
+ params,
+ lr=1e-3,
+ betas=(0.9, 0.98),
+ eps=1e-8,
+ weight_decay=1e-3,
+ target_rms=0.1,
+ ):
+ if not 0.0 <= lr:
+ raise ValueError("Invalid learning rate: {}".format(lr))
+ if not 0.0 <= eps:
+ raise ValueError("Invalid epsilon value: {}".format(eps))
+ if not 0.0 <= betas[0] < 1.0:
+ raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+ if not 0.0 <= betas[1] < 1.0:
+ raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+ if not 0 <= weight_decay <= 0.1:
+ raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+ if not 0 < target_rms <= 10.0:
+ raise ValueError("Invalid target_rms value: {}".format(target_rms))
+ defaults = dict(
+ lr=lr,
+ betas=betas,
+ eps=eps,
+ weight_decay=weight_decay,
+ target_rms=target_rms,
+ )
+ super(Eve, self).__init__(params, defaults)
+
+ def __setstate__(self, state):
+ super(Eve, self).__setstate__(state)
+
+ @torch.no_grad()
+ def step(self, closure=None):
+ """Performs a single optimization step.
+
+ Arguments:
+ closure (callable, optional): A closure that reevaluates the model
+ and returns the loss.
+ """
+ loss = None
+ if closure is not None:
+ with torch.enable_grad():
+ loss = closure()
+
+ for group in self.param_groups:
+ for p in group["params"]:
+ if p.grad is None:
+ continue
+
+ # Perform optimization step
+ grad = p.grad
+ if grad.is_sparse:
+ raise RuntimeError("AdamW does not support sparse gradients")
+
+ state = self.state[p]
+
+ # State initialization
+ if len(state) == 0:
+ state["step"] = 0
+ # Exponential moving average of gradient values
+ state["exp_avg"] = torch.zeros_like(
+ p, memory_format=torch.preserve_format
+ )
+ # Exponential moving average of squared gradient values
+ state["exp_avg_sq"] = torch.zeros_like(
+ p, memory_format=torch.preserve_format
+ )
+
+ exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+ beta1, beta2 = group["betas"]
+
+ state["step"] += 1
+ bias_correction1 = 1 - beta1 ** state["step"]
+ bias_correction2 = 1 - beta2 ** state["step"]
+
+ # Decay the first and second moment running average coefficient
+ exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+ exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+ denom = (exp_avg_sq.sqrt() * (bias_correction2**-0.5)).add_(
+ group["eps"]
+ )
+
+ step_size = group["lr"] / bias_correction1
+ target_rms = group["target_rms"]
+ weight_decay = group["weight_decay"]
+
+ if p.numel() > 1:
+ # avoid applying this weight-decay on "scaling factors"
+ # (which are scalar).
+ is_above_target_rms = p.norm() > (target_rms * (p.numel() ** 0.5))
+ p.mul_(1 - (weight_decay * is_above_target_rms))
+
+ p.addcdiv_(exp_avg, denom, value=-step_size)
+
+ if random.random() < 0.0005:
+ step = (exp_avg / denom) * step_size
+ logging.info(
+ f"Delta rms = {(step**2).mean().item()}, shape = {step.shape}"
+ )
+
+ return loss
+
+
+def _test_scaled_adam(hidden_dim: int):
+ import timeit
+
+ from scaling import ScaledLinear
+
+ E = 100
+ B = 4
+ T = 2
+ logging.info("in test_eve_cain")
+ # device = torch.device('cuda')
+ device = torch.device("cpu")
+ dtype = torch.float32
+
+ fix_random_seed(42)
+ # these input_magnitudes and output_magnitudes are to test that
+ # Abel is working as we expect and is able to adjust scales of
+ # different dims differently.
+ input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+ output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+
+ for iter in [1, 0]:
+ fix_random_seed(42)
+ Linear = torch.nn.Linear if iter == 0 else ScaledLinear
+
+ m = torch.nn.Sequential(
+ Linear(E, hidden_dim),
+ torch.nn.PReLU(),
+ Linear(hidden_dim, hidden_dim),
+ torch.nn.PReLU(),
+ Linear(hidden_dim, E),
+ ).to(device)
+
+ train_pairs = [
+ (
+ 100.0
+ * torch.randn(B, T, E, device=device, dtype=dtype)
+ * input_magnitudes,
+ torch.randn(B, T, E, device=device, dtype=dtype) * output_magnitudes,
+ )
+ for _ in range(20)
+ ]
+
+ if iter == 0:
+ optim = Eve(m.parameters(), lr=0.003)
+ elif iter == 1:
+ optim = ScaledAdam(m.named_parameters(), lr=0.03, clipping_scale=2.0)
+ scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False)
+
+ start = timeit.default_timer()
+ avg_loss = 0.0
+ for epoch in range(180):
+ scheduler.step_epoch()
+ # if epoch == 100 and iter in [2,3]:
+ # optim.reset_speedup() # check it doesn't crash.
+
+ # if epoch == 130:
+ # opts = diagnostics.TensorDiagnosticOptions(
+ # 512
+ # ) # allow 4 megabytes per sub-module
+ # diagnostic = diagnostics.attach_diagnostics(m, opts)
+
+ for n, (x, y) in enumerate(train_pairs):
+ y_out = m(x)
+ loss = ((y_out - y) ** 2).mean() * 100.0
+ if epoch == 0 and n == 0:
+ avg_loss = loss.item()
+ else:
+ avg_loss = 0.98 * avg_loss + 0.02 * loss.item()
+ if n == 0 and epoch % 5 == 0:
+ # norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item()
+ # norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item()
+ # norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item()
+ # norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item()
+ # scale1 = '%.2e' % (m[0].weight_scale.exp().item())
+ # scale1b = '%.2e' % (m[0].bias_scale.exp().item())
+ # scale2 = '%.2e' % (m[2].weight_scale.exp().item())
+ # scale2b = '%.2e' % (m[2].bias_scale.exp().item())
+ lr = scheduler.get_last_lr()[0]
+ logging.info(
+ f"Iter {iter}, epoch {epoch}, batch {n}, avg_loss {avg_loss:.4g}, lr={lr:.4e}"
+ ) # , norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
+ loss.log().backward()
+ optim.step()
+ optim.zero_grad()
+ scheduler.step_batch()
+
+ # diagnostic.print_diagnostics()
+
+ stop = timeit.default_timer()
+ logging.info(f"Iter={iter}, Time taken: {stop - start}")
+
+ logging.info(f"last lr = {scheduler.get_last_lr()}")
+ # logging.info("state dict = ", scheduler.state_dict())
+ # logging.info("optim state_dict = ", optim.state_dict())
+ logging.info(f"input_magnitudes = {input_magnitudes}")
+ logging.info(f"output_magnitudes = {output_magnitudes}")
+
+
+if __name__ == "__main__":
+ torch.set_num_threads(1)
+ torch.set_num_interop_threads(1)
+ logging.getLogger().setLevel(logging.INFO)
+ import subprocess
+
+ s = subprocess.check_output(
+ "git status -uno .; git log -1; git diff HEAD .", shell=True
+ )
+ logging.info(s)
+ import sys
+
+ if len(sys.argv) > 1:
+ hidden_dim = int(sys.argv[1])
+ else:
+ hidden_dim = 200
+
+ _test_scaled_adam(hidden_dim)
+ _test_eden()
diff --git a/egs/europarl_st/SRT/lcma_srt/scaling.py b/egs/europarl_st/SRT/lcma_srt/scaling.py
new file mode 100644
index 0000000000..6d6281903d
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/scaling.py
@@ -0,0 +1,1911 @@
+# Copyright 2022-2023 Xiaomi Corp. (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+import math
+import random
+from typing import Optional, Tuple, Union
+
+import k2
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+
+def logaddexp_onnx(x: Tensor, y: Tensor) -> Tensor:
+ max_value = torch.max(x, y)
+ diff = torch.abs(x - y)
+ return max_value + torch.log1p(torch.exp(-diff))
+
+
+# RuntimeError: Exporting the operator logaddexp to ONNX opset version
+# 14 is not supported. Please feel free to request support or submit
+# a pull request on PyTorch GitHub.
+#
+# The following function is to solve the above error when exporting
+# models to ONNX via torch.jit.trace()
+def logaddexp(x: Tensor, y: Tensor) -> Tensor:
+ # Caution(fangjun): Put torch.jit.is_scripting() before
+ # torch.onnx.is_in_onnx_export();
+ # otherwise, it will cause errors for torch.jit.script().
+ #
+ # torch.logaddexp() works for both torch.jit.script() and
+ # torch.jit.trace() but it causes errors for ONNX export.
+ #
+ if torch.jit.is_scripting():
+ # Note: We cannot use torch.jit.is_tracing() here as it also
+ # matches torch.onnx.export().
+ return torch.logaddexp(x, y)
+ elif torch.onnx.is_in_onnx_export():
+ return logaddexp_onnx(x, y)
+ else:
+ # for torch.jit.trace()
+ return torch.logaddexp(x, y)
+
+
+class PiecewiseLinear(object):
+ """
+ Piecewise linear function, from float to float, specified as nonempty list of (x,y) pairs with
+ the x values in order. x values <[initial x] or >[final x] are map to [initial y], [final y]
+ respectively.
+ """
+
+ def __init__(self, *args):
+ assert len(args) >= 1, len(args)
+ if len(args) == 1 and isinstance(args[0], PiecewiseLinear):
+ self.pairs = list(args[0].pairs)
+ else:
+ self.pairs = [(float(x), float(y)) for x, y in args]
+ for x, y in self.pairs:
+ assert isinstance(x, (float, int)), type(x)
+ assert isinstance(y, (float, int)), type(y)
+
+ for i in range(len(self.pairs) - 1):
+ assert self.pairs[i + 1][0] > self.pairs[i][0], (
+ i,
+ self.pairs[i],
+ self.pairs[i + 1],
+ )
+
+ def __str__(self):
+ # e.g. 'PiecewiseLinear((0., 10.), (100., 0.))'
+ return f"PiecewiseLinear({str(self.pairs)[1:-1]})"
+
+ def __call__(self, x):
+ if x <= self.pairs[0][0]:
+ return self.pairs[0][1]
+ elif x >= self.pairs[-1][0]:
+ return self.pairs[-1][1]
+ else:
+ cur_x, cur_y = self.pairs[0]
+ for i in range(1, len(self.pairs)):
+ next_x, next_y = self.pairs[i]
+ if x >= cur_x and x <= next_x:
+ return cur_y + (next_y - cur_y) * (x - cur_x) / (next_x - cur_x)
+ cur_x, cur_y = next_x, next_y
+ assert False
+
+ def __mul__(self, alpha):
+ return PiecewiseLinear(*[(x, y * alpha) for x, y in self.pairs])
+
+ def __add__(self, x):
+ if isinstance(x, (float, int)):
+ return PiecewiseLinear(*[(p[0], p[1] + x) for p in self.pairs])
+ s, x = self.get_common_basis(x)
+ return PiecewiseLinear(
+ *[(sp[0], sp[1] + xp[1]) for sp, xp in zip(s.pairs, x.pairs)]
+ )
+
+ def max(self, x):
+ if isinstance(x, (float, int)):
+ x = PiecewiseLinear((0, x))
+ s, x = self.get_common_basis(x, include_crossings=True)
+ return PiecewiseLinear(
+ *[(sp[0], max(sp[1], xp[1])) for sp, xp in zip(s.pairs, x.pairs)]
+ )
+
+ def min(self, x):
+ if isinstance(x, float) or isinstance(x, int):
+ x = PiecewiseLinear((0, x))
+ s, x = self.get_common_basis(x, include_crossings=True)
+ return PiecewiseLinear(
+ *[(sp[0], min(sp[1], xp[1])) for sp, xp in zip(s.pairs, x.pairs)]
+ )
+
+ def __eq__(self, other):
+ return self.pairs == other.pairs
+
+ def get_common_basis(self, p: "PiecewiseLinear", include_crossings: bool = False):
+ """
+ Returns (self_mod, p_mod) which are equivalent piecewise linear
+ functions to self and p, but with the same x values.
+
+ p: the other piecewise linear function
+ include_crossings: if true, include in the x values positions
+ where the functions indicate by this and p cross.
+ """
+ assert isinstance(p, PiecewiseLinear), type(p)
+
+ # get sorted x-values without repetition.
+ x_vals = sorted(set([x for x, _ in self.pairs] + [x for x, _ in p.pairs]))
+ y_vals1 = [self(x) for x in x_vals]
+ y_vals2 = [p(x) for x in x_vals]
+
+ if include_crossings:
+ extra_x_vals = []
+ for i in range(len(x_vals) - 1):
+ if (y_vals1[i] > y_vals2[i]) != (y_vals1[i + 1] > y_vals2[i + 1]):
+ # if the two lines in this subsegment potentially cross each other..
+ diff_cur = abs(y_vals1[i] - y_vals2[i])
+ diff_next = abs(y_vals1[i + 1] - y_vals2[i + 1])
+ # `pos`, between 0 and 1, gives the relative x position,
+ # with 0 being x_vals[i] and 1 being x_vals[i+1].
+ pos = diff_cur / (diff_cur + diff_next)
+ extra_x_val = x_vals[i] + pos * (x_vals[i + 1] - x_vals[i])
+ extra_x_vals.append(extra_x_val)
+ if len(extra_x_vals) > 0:
+ x_vals = sorted(set(x_vals + extra_x_vals))
+
+ y_vals1 = [self(x) for x in x_vals]
+ y_vals2 = [p(x) for x in x_vals]
+
+ return (
+ PiecewiseLinear(*zip(x_vals, y_vals1)),
+ PiecewiseLinear(*zip(x_vals, y_vals2)),
+ )
+
+
+class ScheduledFloat(torch.nn.Module):
+ """
+ This object is a torch.nn.Module only because we want it to show up in [top_level module].modules();
+ it does not have a working forward() function. You are supposed to cast it to float, as
+ in, float(parent_module.whatever), and use it as something like a dropout prob.
+
+ It is a floating point value whose value changes depending on the batch count of the
+ training loop. It is a piecewise linear function where you specify the (x,y) pairs
+ in sorted order on x; x corresponds to the batch index. For batch-index values before the
+ first x or after the last x, we just use the first or last y value.
+
+ Example:
+ self.dropout = ScheduledFloat((0.0, 0.2), (4000.0, 0.0), default=0.0)
+
+ `default` is used when self.batch_count is not set or not in training mode or in
+ torch.jit scripting mode.
+ """
+
+ def __init__(self, *args, default: float = 0.0):
+ super().__init__()
+ # self.batch_count and self.name will be written to in the training loop.
+ self.batch_count = None
+ self.name = None
+ self.default = default
+ self.schedule = PiecewiseLinear(*args)
+
+ def extra_repr(self) -> str:
+ return (
+ f"batch_count={self.batch_count}, schedule={str(self.schedule.pairs[1:-1])}"
+ )
+
+ def __float__(self):
+ batch_count = self.batch_count
+ if (
+ batch_count is None
+ or not self.training
+ or torch.jit.is_scripting()
+ or torch.jit.is_tracing()
+ ):
+ return float(self.default)
+ else:
+ ans = self.schedule(self.batch_count)
+ if random.random() < 0.0002:
+ logging.info(
+ f"ScheduledFloat: name={self.name}, batch_count={self.batch_count}, ans={ans}"
+ )
+ return ans
+
+ def __add__(self, x):
+ if isinstance(x, float) or isinstance(x, int):
+ return ScheduledFloat(self.schedule + x, default=self.default)
+ else:
+ return ScheduledFloat(
+ self.schedule + x.schedule, default=self.default + x.default
+ )
+
+ def max(self, x):
+ if isinstance(x, float) or isinstance(x, int):
+ return ScheduledFloat(self.schedule.max(x), default=self.default)
+ else:
+ return ScheduledFloat(
+ self.schedule.max(x.schedule), default=max(self.default, x.default)
+ )
+
+
+FloatLike = Union[float, ScheduledFloat]
+
+
+def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor:
+ """
+ A randomized way of casting a floating point value to half precision.
+ """
+ if x.dtype == torch.float16:
+ return x
+ x_abs = x.abs()
+ is_too_small = x_abs < min_abs
+ # for elements where is_too_small is true, random_val will contain +-min_abs with
+ # probability (x.abs() / min_abs), and 0.0 otherwise. [so this preserves expectations,
+ # for those elements].
+ random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs)
+ return torch.where(is_too_small, random_val, x).to(torch.float16)
+
+
+class CutoffEstimator:
+ """
+ Estimates cutoffs of an arbitrary numerical quantity such that a specified
+ proportion of items will be above the cutoff on average.
+
+ p is the proportion of items that should be above the cutoff.
+ """
+
+ def __init__(self, p: float):
+ self.p = p
+ # total count of items
+ self.count = 0
+ # total count of items that were above the cutoff
+ self.count_above = 0
+ # initial cutoff value
+ self.cutoff = 0
+
+ def __call__(self, x: float) -> bool:
+ """
+ Returns true if x is above the cutoff.
+ """
+ ans = x > self.cutoff
+ self.count += 1
+ if ans:
+ self.count_above += 1
+ cur_p = self.count_above / self.count
+ delta_p = cur_p - self.p
+ if (delta_p > 0) == ans:
+ q = abs(delta_p)
+ self.cutoff = x * q + self.cutoff * (1 - q)
+ return ans
+
+
+class SoftmaxFunction(torch.autograd.Function):
+ """
+ Tries to handle half-precision derivatives in a randomized way that should
+ be more accurate for training than the default behavior.
+ """
+
+ @staticmethod
+ def forward(ctx, x: Tensor, dim: int):
+ ans = x.softmax(dim=dim)
+ # if x dtype is float16, x.softmax() returns a float32 because
+ # (presumably) that op does not support float16, and autocast
+ # is enabled.
+ if torch.is_autocast_enabled():
+ ans = ans.to(torch.get_autocast_gpu_dtype())
+ ctx.save_for_backward(ans)
+ ctx.x_dtype = x.dtype
+ ctx.dim = dim
+ return ans
+
+ @staticmethod
+ def backward(ctx, ans_grad: Tensor):
+ (ans,) = ctx.saved_tensors
+ with torch.cuda.amp.autocast(enabled=False):
+ ans_grad = ans_grad.to(torch.float32)
+ ans = ans.to(torch.float32)
+ x_grad = ans_grad * ans
+ x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True)
+ return x_grad, None
+
+
+def softmax(x: Tensor, dim: int):
+ if not x.requires_grad or torch.jit.is_scripting() or torch.jit.is_tracing():
+ return x.softmax(dim=dim)
+
+ return SoftmaxFunction.apply(x, dim)
+
+
+class MaxEigLimiterFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx,
+ x: Tensor,
+ coeffs: Tensor,
+ direction: Tensor,
+ channel_dim: int,
+ grad_scale: float,
+ ) -> Tensor:
+ ctx.channel_dim = channel_dim
+ ctx.grad_scale = grad_scale
+ ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach())
+ return x
+
+ @staticmethod
+ def backward(ctx, x_grad, *args):
+ with torch.enable_grad():
+ (x_orig, coeffs, new_direction) = ctx.saved_tensors
+ x_orig.requires_grad = True
+ num_channels = x_orig.shape[ctx.channel_dim]
+ x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels)
+ new_direction.requires_grad = False
+ x = x - x.mean(dim=0)
+ x_var = (x**2).mean()
+ x_residual = x - coeffs * new_direction
+ x_residual_var = (x_residual**2).mean()
+ # `variance_proportion` is the proportion of the variance accounted for
+ # by the top eigen-direction. This is to be minimized.
+ variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+ variance_proportion.backward()
+ x_orig_grad = x_orig.grad
+ x_extra_grad = (
+ x_orig.grad
+ * ctx.grad_scale
+ * x_grad.norm()
+ / (x_orig_grad.norm() + 1.0e-20)
+ )
+ return x_grad + x_extra_grad.detach(), None, None, None, None
+
+
+class BiasNormFunction(torch.autograd.Function):
+ # This computes:
+ # scales = (torch.mean((x - bias) ** 2, keepdim=True)) ** -0.5 * log_scale.exp()
+ # return x * scales
+ # (after unsqueezing the bias), but it does it in a memory-efficient way so that
+ # it can just store the returned value (chances are, this will also be needed for
+ # some other reason, related to the next operation, so we can save memory).
+ @staticmethod
+ def forward(
+ ctx,
+ x: Tensor,
+ bias: Tensor,
+ log_scale: Tensor,
+ channel_dim: int,
+ store_output_for_backprop: bool,
+ ) -> Tensor:
+ assert bias.ndim == 1
+ if channel_dim < 0:
+ channel_dim = channel_dim + x.ndim
+ ctx.store_output_for_backprop = store_output_for_backprop
+ ctx.channel_dim = channel_dim
+ for _ in range(channel_dim + 1, x.ndim):
+ bias = bias.unsqueeze(-1)
+ scales = (
+ torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5
+ ) * log_scale.exp()
+ ans = x * scales
+ ctx.save_for_backward(
+ ans.detach() if store_output_for_backprop else x,
+ scales.detach(),
+ bias.detach(),
+ log_scale.detach(),
+ )
+ return ans
+
+ @staticmethod
+ def backward(ctx, ans_grad: Tensor) -> Tensor:
+ ans_or_x, scales, bias, log_scale = ctx.saved_tensors
+ if ctx.store_output_for_backprop:
+ x = ans_or_x / scales
+ else:
+ x = ans_or_x
+ x = x.detach()
+ x.requires_grad = True
+ bias.requires_grad = True
+ log_scale.requires_grad = True
+ with torch.enable_grad():
+ # recompute scales from x, bias and log_scale.
+ scales = (
+ torch.mean((x - bias) ** 2, dim=ctx.channel_dim, keepdim=True) ** -0.5
+ ) * log_scale.exp()
+ ans = x * scales
+ ans.backward(gradient=ans_grad)
+ return x.grad, bias.grad.flatten(), log_scale.grad, None, None
+
+
+class BiasNorm(torch.nn.Module):
+ """
+ This is intended to be a simpler, and hopefully cheaper, replacement for
+ LayerNorm. The observation this is based on, is that Transformer-type
+ networks, especially with pre-norm, sometimes seem to set one of the
+ feature dimensions to a large constant value (e.g. 50), which "defeats"
+ the LayerNorm because the output magnitude is then not strongly dependent
+ on the other (useful) features. Presumably the weight and bias of the
+ LayerNorm are required to allow it to do this.
+
+ Instead, we give the BiasNorm a trainable bias that it can use when
+ computing the scale for normalization. We also give it a (scalar)
+ trainable scale on the output.
+
+
+ Args:
+ num_channels: the number of channels, e.g. 512.
+ channel_dim: the axis/dimension corresponding to the channel,
+ interpreted as an offset from the input's ndim if negative.
+ This is NOT the num_channels; it should typically be one of
+ {-2, -1, 0, 1, 2, 3}.
+ log_scale: the initial log-scale that we multiply the output by; this
+ is learnable.
+ log_scale_min: FloatLike, minimum allowed value of log_scale
+ log_scale_max: FloatLike, maximum allowed value of log_scale
+ store_output_for_backprop: only possibly affects memory use; recommend
+ to set to True if you think the output of this module is more likely
+ than the input of this module to be required to be stored for the
+ backprop.
+ """
+
+ def __init__(
+ self,
+ num_channels: int,
+ channel_dim: int = -1, # CAUTION: see documentation.
+ log_scale: float = 1.0,
+ log_scale_min: float = -1.5,
+ log_scale_max: float = 1.5,
+ store_output_for_backprop: bool = False,
+ ) -> None:
+ super(BiasNorm, self).__init__()
+ self.num_channels = num_channels
+ self.channel_dim = channel_dim
+ self.log_scale = nn.Parameter(torch.tensor(log_scale))
+ self.bias = nn.Parameter(torch.empty(num_channels).normal_(mean=0, std=1e-4))
+
+ self.log_scale_min = log_scale_min
+ self.log_scale_max = log_scale_max
+
+ self.store_output_for_backprop = store_output_for_backprop
+
+ def forward(self, x: Tensor) -> Tensor:
+ assert x.shape[self.channel_dim] == self.num_channels
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ channel_dim = self.channel_dim
+ if channel_dim < 0:
+ channel_dim += x.ndim
+ bias = self.bias
+ for _ in range(channel_dim + 1, x.ndim):
+ bias = bias.unsqueeze(-1)
+ scales = (
+ torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5
+ ) * self.log_scale.exp()
+ return x * scales
+
+ log_scale = limit_param_value(
+ self.log_scale,
+ min=float(self.log_scale_min),
+ max=float(self.log_scale_max),
+ training=self.training,
+ )
+
+ return BiasNormFunction.apply(
+ x, self.bias, log_scale, self.channel_dim, self.store_output_for_backprop
+ )
+
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+ """
+ Behaves like a constructor of a modified version of nn.Linear
+ that gives an easy way to set the default initial parameter scale.
+
+ Args:
+ Accepts the standard args and kwargs that nn.Linear accepts
+ e.g. in_features, out_features, bias=False.
+
+ initial_scale: you can override this if you want to increase
+ or decrease the initial magnitude of the module's output
+ (affects the initialization of weight_scale and bias_scale).
+ Another option, if you want to do something like this, is
+ to re-initialize the parameters.
+ """
+ ans = nn.Linear(*args, **kwargs)
+ with torch.no_grad():
+ ans.weight[:] *= initial_scale
+ if ans.bias is not None:
+ torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale)
+ return ans
+
+
+def ScaledConv1d(*args, initial_scale: float = 1.0, **kwargs) -> nn.Conv1d:
+ """
+ Behaves like a constructor of a modified version of nn.Conv1d
+ that gives an easy way to set the default initial parameter scale.
+
+ Args:
+ Accepts the standard args and kwargs that nn.Linear accepts
+ e.g. in_features, out_features, bias=False.
+
+ initial_scale: you can override this if you want to increase
+ or decrease the initial magnitude of the module's output
+ (affects the initialization of weight_scale and bias_scale).
+ Another option, if you want to do something like this, is
+ to re-initialize the parameters.
+ """
+ ans = nn.Conv1d(*args, **kwargs)
+ with torch.no_grad():
+ ans.weight[:] *= initial_scale
+ if ans.bias is not None:
+ torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale)
+ return ans
+
+
+def ScaledConv2d(*args, initial_scale: float = 1.0, **kwargs) -> nn.Conv2d:
+ """
+ Behaves like a constructor of a modified version of nn.Conv2d
+ that gives an easy way to set the default initial parameter scale.
+
+ Args:
+ Accepts the standard args and kwargs that nn.Linear accepts
+ e.g. in_features, out_features, bias=False, but:
+ NO PADDING-RELATED ARGS.
+
+ initial_scale: you can override this if you want to increase
+ or decrease the initial magnitude of the module's output
+ (affects the initialization of weight_scale and bias_scale).
+ Another option, if you want to do something like this, is
+ to re-initialize the parameters.
+ """
+ ans = nn.Conv2d(*args, **kwargs)
+ with torch.no_grad():
+ ans.weight[:] *= initial_scale
+ if ans.bias is not None:
+ torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale)
+ return ans
+
+
+class ChunkCausalDepthwiseConv1d(torch.nn.Module):
+ """
+ Behaves like a depthwise 1d convolution, except that it is causal in
+ a chunkwise way, as if we had a block-triangular attention mask.
+ The chunk size is provided at test time (it should probably be
+ kept in sync with the attention mask).
+
+ This has a little more than twice the parameters of a conventional
+ depthwise conv1d module: we implement it by having one
+ depthwise convolution, of half the width, that is causal (via
+ right-padding); and one depthwise convolution that is applied only
+ within chunks, that we multiply by a scaling factor which depends
+ on the position within the chunk.
+
+ Args:
+ Accepts the standard args and kwargs that nn.Linear accepts
+ e.g. in_features, out_features, bias=False.
+
+ initial_scale: you can override this if you want to increase
+ or decrease the initial magnitude of the module's output
+ (affects the initialization of weight_scale and bias_scale).
+ Another option, if you want to do something like this, is
+ to re-initialize the parameters.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ kernel_size: int,
+ initial_scale: float = 1.0,
+ bias: bool = True,
+ ):
+ super().__init__()
+ assert kernel_size % 2 == 1
+
+ half_kernel_size = (kernel_size + 1) // 2
+ # will pad manually, on one side.
+ self.causal_conv = nn.Conv1d(
+ in_channels=channels,
+ out_channels=channels,
+ groups=channels,
+ kernel_size=half_kernel_size,
+ padding=0,
+ bias=True,
+ )
+
+ self.chunkwise_conv = nn.Conv1d(
+ in_channels=channels,
+ out_channels=channels,
+ groups=channels,
+ kernel_size=kernel_size,
+ padding=kernel_size // 2,
+ bias=bias,
+ )
+
+ # first row is correction factors added to the scale near the left edge of the chunk,
+ # second row is correction factors added to the scale near the right edge of the chunk,
+ # both of these are added to a default scale of 1.0.
+ self.chunkwise_conv_scale = nn.Parameter(torch.zeros(2, channels, kernel_size))
+ self.kernel_size = kernel_size
+
+ with torch.no_grad():
+ self.causal_conv.weight[:] *= initial_scale
+ self.chunkwise_conv.weight[:] *= initial_scale
+ if bias:
+ torch.nn.init.uniform_(
+ self.causal_conv.bias, -0.1 * initial_scale, 0.1 * initial_scale
+ )
+
+ def forward(self, x: Tensor, chunk_size: int = -1) -> Tensor:
+ """Forward function.
+
+ Args:
+ x: a Tensor of shape (batch_size, channels, seq_len)
+ chunk_size: the chunk size, in frames; does not have to divide seq_len exactly.
+ """
+ (batch_size, num_channels, seq_len) = x.shape
+
+ # half_kernel_size = self.kernel_size + 1 // 2
+ # left_pad is half_kernel_size - 1 where half_kernel_size is the size used
+ # in the causal conv. It's the amount by which we must pad on the left,
+ # to make the convolution causal.
+ left_pad = self.kernel_size // 2
+
+ if chunk_size < 0 or chunk_size > seq_len:
+ chunk_size = seq_len
+ right_pad = -seq_len % chunk_size
+
+ x = torch.nn.functional.pad(x, (left_pad, right_pad))
+
+ x_causal = self.causal_conv(x[..., : left_pad + seq_len])
+ assert x_causal.shape == (batch_size, num_channels, seq_len)
+
+ x_chunk = x[..., left_pad:]
+ num_chunks = x_chunk.shape[2] // chunk_size
+ x_chunk = x_chunk.reshape(batch_size, num_channels, num_chunks, chunk_size)
+ x_chunk = x_chunk.permute(0, 2, 1, 3).reshape(
+ batch_size * num_chunks, num_channels, chunk_size
+ )
+ x_chunk = self.chunkwise_conv(x_chunk) # does not change shape
+
+ chunk_scale = self._get_chunk_scale(chunk_size)
+
+ x_chunk = x_chunk * chunk_scale
+ x_chunk = x_chunk.reshape(
+ batch_size, num_chunks, num_channels, chunk_size
+ ).permute(0, 2, 1, 3)
+ x_chunk = x_chunk.reshape(batch_size, num_channels, num_chunks * chunk_size)[
+ ..., :seq_len
+ ]
+
+ return x_chunk + x_causal
+
+ def _get_chunk_scale(self, chunk_size: int):
+ """Returns tensor of shape (num_channels, chunk_size) that will be used to
+ scale the output of self.chunkwise_conv."""
+ left_edge = self.chunkwise_conv_scale[0]
+ right_edge = self.chunkwise_conv_scale[1]
+ if chunk_size < self.kernel_size:
+ left_edge = left_edge[:, :chunk_size]
+ right_edge = right_edge[:, -chunk_size:]
+ else:
+ t = chunk_size - self.kernel_size
+ channels = left_edge.shape[0]
+ pad = torch.zeros(
+ channels, t, device=left_edge.device, dtype=left_edge.dtype
+ )
+ left_edge = torch.cat((left_edge, pad), dim=-1)
+ right_edge = torch.cat((pad, right_edge), dim=-1)
+ return 1.0 + (left_edge + right_edge)
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ cache: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ """Streaming Forward function.
+
+ Args:
+ x: a Tensor of shape (batch_size, channels, seq_len)
+ cache: cached left context of shape (batch_size, channels, left_pad)
+ """
+ (batch_size, num_channels, seq_len) = x.shape
+
+ # left_pad is half_kernel_size - 1 where half_kernel_size is the size used
+ # in the causal conv. It's the amount by which we must pad on the left,
+ # to make the convolution causal.
+ left_pad = self.kernel_size // 2
+
+ # Pad cache
+ assert cache.shape[-1] == left_pad, (cache.shape[-1], left_pad)
+ x = torch.cat([cache, x], dim=2)
+ # Update cache
+ cache = x[..., -left_pad:]
+
+ x_causal = self.causal_conv(x)
+ assert x_causal.shape == (batch_size, num_channels, seq_len)
+
+ x_chunk = x[..., left_pad:]
+ x_chunk = self.chunkwise_conv(x_chunk) # does not change shape
+
+ chunk_scale = self._get_chunk_scale(chunk_size=seq_len)
+ x_chunk = x_chunk * chunk_scale
+
+ return x_chunk + x_causal, cache
+
+
+class BalancerFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx,
+ x: Tensor,
+ min_mean: float,
+ max_mean: float,
+ min_rms: float,
+ max_rms: float,
+ grad_scale: float,
+ channel_dim: int,
+ ) -> Tensor:
+ if channel_dim < 0:
+ channel_dim += x.ndim
+ ctx.channel_dim = channel_dim
+ ctx.save_for_backward(x)
+ ctx.config = (min_mean, max_mean, min_rms, max_rms, grad_scale, channel_dim)
+ return x
+
+ @staticmethod
+ def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None, None, None]:
+ (x,) = ctx.saved_tensors
+ (min_mean, max_mean, min_rms, max_rms, grad_scale, channel_dim) = ctx.config
+
+ try:
+ with torch.enable_grad():
+ with torch.cuda.amp.autocast(enabled=False):
+ x = x.to(torch.float32)
+ x = x.detach()
+ x.requires_grad = True
+ mean_dims = [i for i in range(x.ndim) if i != channel_dim]
+ uncentered_var = (x**2).mean(dim=mean_dims, keepdim=True)
+ mean = x.mean(dim=mean_dims, keepdim=True)
+ stddev = (uncentered_var - (mean * mean)).clamp(min=1.0e-20).sqrt()
+ rms = uncentered_var.clamp(min=1.0e-20).sqrt()
+
+ m = mean / stddev
+ # part of loss that relates to mean / stddev
+ m_loss = (m - m.clamp(min=min_mean, max=max_mean)).abs()
+
+ # put a much larger scale on the RMS-max-limit loss, so that if both it and the
+ # m_loss are violated we fix the RMS loss first.
+ rms_clamped = rms.clamp(min=min_rms, max=max_rms)
+ r_loss = (rms_clamped / rms).log().abs()
+
+ loss = m_loss + r_loss
+
+ loss.backward(gradient=torch.ones_like(loss))
+ loss_grad = x.grad
+ loss_grad_rms = (
+ (loss_grad**2)
+ .mean(dim=mean_dims, keepdim=True)
+ .sqrt()
+ .clamp(min=1.0e-20)
+ )
+
+ loss_grad = loss_grad * (grad_scale / loss_grad_rms)
+
+ x_grad_float = x_grad.to(torch.float32)
+ # scale each element of loss_grad by the absolute value of the corresponding
+ # element of x_grad, which we view as a noisy estimate of its magnitude for that
+ # (frame and dimension). later we can consider factored versions.
+ x_grad_mod = x_grad_float + (x_grad_float.abs() * loss_grad)
+ x_grad = x_grad_mod.to(x_grad.dtype)
+ except Exception as e:
+ logging.info(
+ f"Caught exception in Balancer backward: {e}, size={list(x_grad.shape)}, will continue."
+ )
+
+ return x_grad, None, None, None, None, None, None
+
+
+class Balancer(torch.nn.Module):
+ """
+ Modifies the backpropped derivatives of a function to try to encourage, for
+ each channel, that it is positive at least a proportion `threshold` of the
+ time. It does this by multiplying negative derivative values by up to
+ (1+max_factor), and positive derivative values by up to (1-max_factor),
+ interpolated from 1 at the threshold to those extremal values when none
+ of the inputs are positive.
+
+ Args:
+ num_channels: the number of channels
+ channel_dim: the dimension/axis corresponding to the channel, e.g.
+ -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+ min_positive: the minimum, per channel, of the proportion of the time
+ that (x > 0), below which we start to modify the derivatives.
+ max_positive: the maximum, per channel, of the proportion of the time
+ that (x > 0), above which we start to modify the derivatives.
+ scale_gain_factor: determines the 'gain' with which we increase the
+ change in gradient once the constraints on min_abs and max_abs
+ are violated.
+ min_abs: the minimum average-absolute-value difference from the mean
+ value per channel, which we allow, before we start to modify
+ the derivatives to prevent this.
+ max_abs: the maximum average-absolute-value difference from the mean
+ value per channel, which we allow, before we start to modify
+ the derivatives to prevent this.
+ prob: determines the minimum probability with which we modify the
+ gradients for the {min,max}_positive and {min,max}_abs constraints,
+ on each forward(). This is done randomly to prevent all layers
+ from doing it at the same time.
+ """
+
+ def __init__(
+ self,
+ num_channels: int,
+ channel_dim: int,
+ min_positive: FloatLike = 0.05,
+ max_positive: FloatLike = 0.95,
+ min_abs: FloatLike = 0.2,
+ max_abs: FloatLike = 100.0,
+ grad_scale: FloatLike = 0.04,
+ prob: Optional[FloatLike] = None,
+ ):
+ super().__init__()
+
+ if prob is None:
+ prob = ScheduledFloat((0.0, 0.5), (8000.0, 0.125), default=0.4)
+ self.prob = prob
+ # 5% of the time we will return and do nothing because memory usage is
+ # too high.
+ self.mem_cutoff = CutoffEstimator(0.05)
+
+ # actually self.num_channels is no longer needed except for an assertion.
+ self.num_channels = num_channels
+ self.channel_dim = channel_dim
+ self.min_positive = min_positive
+ self.max_positive = max_positive
+ self.min_abs = min_abs
+ self.max_abs = max_abs
+ self.grad_scale = grad_scale
+
+ def forward(self, x: Tensor) -> Tensor:
+ if (
+ torch.jit.is_scripting()
+ or not x.requires_grad
+ or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated()))
+ ):
+ return _no_op(x)
+
+ prob = float(self.prob)
+ if random.random() < prob:
+ # The following inner-functions convert from the way we historically specified
+ # these limitations, as limits on the absolute value and the proportion of positive
+ # values, to limits on the RMS value and the (mean / stddev).
+ def _abs_to_rms(x):
+ # for normally distributed data, if the expected absolute value is x, the
+ # expected rms value will be sqrt(pi/2) * x.
+ return 1.25331413732 * x
+
+ def _proportion_positive_to_mean(x):
+ def _atanh(x):
+ eps = 1.0e-10
+ # eps is to prevent crashes if x is exactly 0 or 1.
+ # we'll just end up returning a fairly large value.
+ return (math.log(1 + x + eps) - math.log(1 - x + eps)) / 2.0
+
+ def _approx_inverse_erf(x):
+ # 1 / (sqrt(pi) * ln(2)),
+ # see https://math.stackexchange.com/questions/321569/approximating-the-error-function-erf-by-analytical-functions
+ # this approximation is extremely crude and gets progressively worse for
+ # x very close to -1 or +1, but we mostly care about the "middle" region
+ # e.g. _approx_inverse_erf(0.05) = 0.0407316414078772,
+ # and math.erf(0.0407316414078772) = 0.045935330944660666,
+ # which is pretty close to 0.05.
+ return 0.8139535143 * _atanh(x)
+
+ # first convert x from the range 0..1 to the range -1..1 which the error
+ # function returns
+ x = -1 + (2 * x)
+ return _approx_inverse_erf(x)
+
+ min_mean = _proportion_positive_to_mean(float(self.min_positive))
+ max_mean = _proportion_positive_to_mean(float(self.max_positive))
+ min_rms = _abs_to_rms(float(self.min_abs))
+ max_rms = _abs_to_rms(float(self.max_abs))
+ grad_scale = float(self.grad_scale)
+
+ assert x.shape[self.channel_dim] == self.num_channels
+
+ return BalancerFunction.apply(
+ x, min_mean, max_mean, min_rms, max_rms, grad_scale, self.channel_dim
+ )
+ else:
+ return _no_op(x)
+
+
+def penalize_abs_values_gt(
+ x: Tensor, limit: float, penalty: float, name: str = None
+) -> Tensor:
+ """
+ Returns x unmodified, but in backprop will put a penalty for the excess of
+ the absolute values of elements of x over the limit "limit". E.g. if
+ limit == 10.0, then if x has any values over 10 it will get a penalty.
+
+ Caution: the value of this penalty will be affected by grad scaling used
+ in automatic mixed precision training. For this reasons we use this,
+ it shouldn't really matter, or may even be helpful; we just use this
+ to disallow really implausible values of scores to be given to softmax.
+
+ The name is for randomly printed debug info.
+ """
+ x_sign = x.sign()
+ over_limit = (x.abs() - limit) > 0
+ # The following is a memory efficient way to penalize the absolute values of
+ # x that's over the limit. (The memory efficiency comes when you think
+ # about which items torch needs to cache for the autograd, and which ones it
+ # can throw away). The numerical value of aux_loss as computed here will
+ # actually be larger than it should be, by limit * over_limit.sum(), but it
+ # has the same derivative as the real aux_loss which is penalty * (x.abs() -
+ # limit).relu().
+ aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x)
+ # note: we don't do sum() here on aux)_loss, but it's as if we had done
+ # sum() due to how with_loss() works.
+ x = with_loss(x, aux_loss, name)
+ # you must use x for something, or this will be ineffective.
+ return x
+
+
+def _diag(x: Tensor): # like .diag(), but works for tensors with 3 dims.
+ if x.ndim == 2:
+ return x.diag()
+ else:
+ (batch, dim, dim) = x.shape
+ x = x.reshape(batch, dim * dim)
+ x = x[:, :: dim + 1]
+ assert x.shape == (batch, dim)
+ return x
+
+
+def _whitening_metric(x: Tensor, num_groups: int):
+ """
+ Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of
+ of the centered feature covariance are the same within each group's covariance matrix
+ and also between groups.
+ Args:
+ x: a Tensor of shape (*, num_channels)
+ num_groups: the number of groups of channels, a number >=1 that divides num_channels
+ Returns:
+ Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and
+ greater than 1.0 otherwise.
+ """
+ assert x.dtype != torch.float16
+ x = x.reshape(-1, x.shape[-1])
+ (num_frames, num_channels) = x.shape
+ assert num_channels % num_groups == 0
+ channels_per_group = num_channels // num_groups
+ x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1)
+ # x now has shape (num_groups, num_frames, channels_per_group)
+ # subtract the mean so we use the centered, not uncentered, covariance.
+ # My experience has been that when we "mess with the gradients" like this,
+ # it's better not do anything that tries to move the mean around, because
+ # that can easily cause instability.
+ x = x - x.mean(dim=1, keepdim=True)
+ # x_covar: (num_groups, channels_per_group, channels_per_group)
+ x_covar = torch.matmul(x.transpose(1, 2), x)
+ x_covar_mean_diag = _diag(x_covar).mean()
+ # the following expression is what we'd get if we took the matrix product
+ # of each covariance and measured the mean of its trace, i.e.
+ # the same as _diag(torch.matmul(x_covar, x_covar)).mean().
+ x_covarsq_mean_diag = (x_covar**2).sum() / (num_groups * channels_per_group)
+ # this metric will be >= 1.0; the larger it is, the less 'white' the data was.
+ metric = x_covarsq_mean_diag / (x_covar_mean_diag**2 + 1.0e-20)
+ return metric
+
+
+class WhiteningPenaltyFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x: Tensor, module: nn.Module) -> Tensor:
+ ctx.save_for_backward(x)
+ ctx.module = module
+ return x
+
+ @staticmethod
+ def backward(ctx, x_grad: Tensor):
+ (x_orig,) = ctx.saved_tensors
+ w = ctx.module
+
+ try:
+ with torch.enable_grad():
+ with torch.cuda.amp.autocast(enabled=False):
+ x_detached = x_orig.to(torch.float32).detach()
+ x_detached.requires_grad = True
+
+ metric = _whitening_metric(x_detached, w.num_groups)
+
+ if random.random() < 0.005 or __name__ == "__main__":
+ logging.info(
+ f"Whitening: name={w.name}, num_groups={w.num_groups}, num_channels={x_orig.shape[-1]}, "
+ f"metric={metric.item():.2f} vs. limit={float(w.whitening_limit)}"
+ )
+
+ if metric < float(w.whitening_limit):
+ w.prob = w.min_prob
+ return x_grad, None
+ else:
+ w.prob = w.max_prob
+ metric.backward()
+ penalty_grad = x_detached.grad
+ scale = float(w.grad_scale) * (
+ x_grad.to(torch.float32).norm()
+ / (penalty_grad.norm() + 1.0e-20)
+ )
+ penalty_grad = penalty_grad * scale
+ return x_grad + penalty_grad.to(x_grad.dtype), None
+ except Exception as e:
+ logging.info(
+ f"Caught exception in Whiten backward: {e}, size={list(x_grad.shape)}, will continue."
+ )
+ return x_grad, None
+
+
+class Whiten(nn.Module):
+ def __init__(
+ self,
+ num_groups: int,
+ whitening_limit: FloatLike,
+ prob: Union[float, Tuple[float, float]],
+ grad_scale: FloatLike,
+ ):
+ """
+ Args:
+ num_groups: the number of groups to divide the channel dim into before
+ whitening. We will attempt to make the feature covariance
+ within each group, after mean subtraction, as "white" as possible,
+ while having the same trace across all groups.
+ whitening_limit: a value greater than 1.0, that dictates how much
+ freedom we have to violate the constraints. 1.0 would mean perfectly
+ white, with exactly the same trace across groups; larger values
+ give more freedom. E.g. 2.0.
+ prob: the probability with which we apply the gradient modification
+ (also affects the grad scale). May be supplied as a float,
+ or as a pair (min_prob, max_prob)
+
+ grad_scale: determines the scale on the gradient term from this object,
+ relative to the rest of the gradient on the attention weights.
+ E.g. 0.02 (you may want to use smaller values than this if prob is large)
+ """
+ super(Whiten, self).__init__()
+ assert num_groups >= 1
+ assert float(whitening_limit) >= 1
+ assert float(grad_scale) >= 0
+ self.num_groups = num_groups
+ self.whitening_limit = whitening_limit
+ self.grad_scale = grad_scale
+
+ if isinstance(prob, float):
+ prob = (prob, prob)
+ (self.min_prob, self.max_prob) = prob
+ assert 0 < self.min_prob <= self.max_prob <= 1
+ self.prob = self.max_prob
+ self.name = None # will be set in training loop
+
+ def forward(self, x: Tensor) -> Tensor:
+ """
+ In the forward pass, this function just returns the input unmodified.
+ In the backward pass, it will modify the gradients to ensure that the
+ distribution in each group has close to (lambda times I) as the covariance
+ after mean subtraction, with the same lambda across groups.
+ For whitening_limit > 1, there will be more freedom to violate this
+ constraint.
+
+ Args:
+ x: the input of shape (*, num_channels)
+
+ Returns:
+ x, unmodified. You should make sure
+ you use the returned value, or the graph will be freed
+ and nothing will happen in backprop.
+ """
+ grad_scale = float(self.grad_scale)
+ if not x.requires_grad or random.random() > self.prob or grad_scale == 0:
+ return _no_op(x)
+ else:
+ return WhiteningPenaltyFunction.apply(x, self)
+
+
+class WithLoss(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x: Tensor, y: Tensor, name: str):
+ ctx.y_shape = y.shape
+ if random.random() < 0.002 and name is not None:
+ loss_sum = y.sum().item()
+ logging.info(f"WithLoss: name={name}, loss-sum={loss_sum:.3e}")
+ return x
+
+ @staticmethod
+ def backward(ctx, ans_grad: Tensor):
+ return (
+ ans_grad,
+ torch.ones(ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device),
+ None,
+ )
+
+
+def with_loss(x, y, name):
+ # returns x but adds y.sum() to the loss function.
+ return WithLoss.apply(x, y, name)
+
+
+class ScaleGradFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x: Tensor, alpha: float) -> Tensor:
+ ctx.alpha = alpha
+ return x
+
+ @staticmethod
+ def backward(ctx, grad: Tensor):
+ return grad * ctx.alpha, None
+
+
+def scale_grad(x: Tensor, alpha: float):
+ return ScaleGradFunction.apply(x, alpha)
+
+
+class ScaleGrad(nn.Module):
+ def __init__(self, alpha: float):
+ super().__init__()
+ self.alpha = alpha
+
+ def forward(self, x: Tensor) -> Tensor:
+ if torch.jit.is_scripting() or torch.jit.is_tracing() or not self.training:
+ return x
+ return scale_grad(x, self.alpha)
+
+
+class LimitParamValue(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x: Tensor, min: float, max: float):
+ ctx.save_for_backward(x)
+ assert max >= min
+ ctx.min = min
+ ctx.max = max
+ return x
+
+ @staticmethod
+ def backward(ctx, x_grad: Tensor):
+ (x,) = ctx.saved_tensors
+ # where x < ctx.min, ensure all grads are negative (this will tend to make
+ # x more positive).
+ x_grad = x_grad * torch.where(
+ torch.logical_and(x_grad > 0, x < ctx.min), -1.0, 1.0
+ )
+ # where x > ctx.max, ensure all grads are positive (this will tend to make
+ # x more negative).
+ x_grad *= torch.where(torch.logical_and(x_grad < 0, x > ctx.max), -1.0, 1.0)
+ return x_grad, None, None
+
+
+def limit_param_value(
+ x: Tensor, min: float, max: float, prob: float = 0.6, training: bool = True
+):
+ # You apply this to (typically) an nn.Parameter during training to ensure that its
+ # (elements mostly) stays within a supplied range. This is done by modifying the
+ # gradients in backprop.
+ # It's not necessary to do this on every batch: do it only some of the time,
+ # to save a little time.
+ if training and random.random() < prob:
+ return LimitParamValue.apply(x, min, max)
+ else:
+ return x
+
+
+def _no_op(x: Tensor) -> Tensor:
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ return x
+ else:
+ # a no-op function that will have a node in the autograd graph,
+ # to avoid certain bugs relating to backward hooks
+ return x.chunk(1, dim=-1)[0]
+
+
+class Identity(torch.nn.Module):
+ def __init__(self):
+ super(Identity, self).__init__()
+
+ def forward(self, x):
+ return _no_op(x)
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+ """
+ double_swish(x) = x * torch.sigmoid(x-1)
+
+ This is a definition, originally motivated by its close numerical
+ similarity to swish(swish(x)), where swish(x) = x * sigmoid(x).
+
+ Memory-efficient derivative computation:
+ double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+ double_swish'(x) = d/dx double_swish(x) = x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+ Now, s'(x) = s(x) * (1-s(x)).
+ double_swish'(x) = x * s'(x) + s(x).
+ = x * s(x) * (1-s(x)) + s(x).
+ = double_swish(x) * (1-s(x)) + s(x)
+ ... so we just need to remember s(x) but not x itself.
+ """
+
+ @staticmethod
+ def forward(ctx, x: Tensor) -> Tensor:
+ requires_grad = x.requires_grad
+ if x.dtype == torch.float16 or x.dtype == torch.bfloat16:
+ x = x.to(torch.float32)
+
+ s = torch.sigmoid(x - 1.0)
+ y = x * s
+
+ if requires_grad:
+ deriv = y * (1 - s) + s
+
+ # notes on derivative of x * sigmoid(x - 1):
+ # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+ # min \simeq -0.043638. Take floor as -0.044 so it's a lower bund
+ # max \simeq 1.1990. Take ceil to be 1.2 so it's an upper bound.
+ # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+ # floors), should be expectation-preserving.
+ floor = -0.044
+ ceil = 1.2
+ d_scaled = (deriv - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
+ deriv
+ )
+ if __name__ == "__main__":
+ # for self-testing only.
+ assert d_scaled.min() >= 0.0
+ assert d_scaled.max() < 256.0
+ d_int = d_scaled.to(torch.uint8)
+ ctx.save_for_backward(d_int)
+ if x.dtype == torch.float16 or torch.is_autocast_enabled():
+ y = y.to(torch.float16)
+ return y
+
+ @staticmethod
+ def backward(ctx, y_grad: Tensor) -> Tensor:
+ (d,) = ctx.saved_tensors
+ # the same constants as used in forward pass.
+ floor = -0.043637
+ ceil = 1.2
+
+ d = d * ((ceil - floor) / 255.0) + floor
+ return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, x: Tensor) -> Tensor:
+ """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+ that we approximate closely with x * sigmoid(x-1).
+ """
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ return x * torch.sigmoid(x - 1.0)
+ return DoubleSwishFunction.apply(x)
+
+
+# Dropout2 is just like normal dropout, except it supports schedules on the dropout rates.
+class Dropout2(nn.Module):
+ def __init__(self, p: FloatLike):
+ super().__init__()
+ self.p = p
+
+ def forward(self, x: Tensor) -> Tensor:
+ return torch.nn.functional.dropout(x, p=float(self.p), training=self.training)
+
+
+class MulForDropout3(torch.autograd.Function):
+ # returns (x * y * alpha) where alpha is a float and y doesn't require
+ # grad and is zero-or-one.
+ @staticmethod
+ @custom_fwd
+ def forward(ctx, x, y, alpha):
+ assert not y.requires_grad
+ ans = x * y * alpha
+ ctx.save_for_backward(ans)
+ ctx.alpha = alpha
+ return ans
+
+ @staticmethod
+ @custom_bwd
+ def backward(ctx, ans_grad):
+ (ans,) = ctx.saved_tensors
+ x_grad = ctx.alpha * ans_grad * (ans != 0)
+ return x_grad, None, None
+
+
+# Dropout3 is just like normal dropout, except it supports schedules on the dropout rates,
+# and it lets you choose one dimension to share the dropout mask over
+class Dropout3(nn.Module):
+ def __init__(self, p: FloatLike, shared_dim: int):
+ super().__init__()
+ self.p = p
+ self.shared_dim = shared_dim
+
+ def forward(self, x: Tensor) -> Tensor:
+ p = float(self.p)
+ if not self.training or p == 0:
+ return _no_op(x)
+ scale = 1.0 / (1 - p)
+ rand_shape = list(x.shape)
+ rand_shape[self.shared_dim] = 1
+ mask = torch.rand(*rand_shape, device=x.device) > p
+ ans = MulForDropout3.apply(x, mask, scale)
+ return ans
+
+
+class SwooshLFunction(torch.autograd.Function):
+ """
+ swoosh_l(x) = log(1 + exp(x-4)) - 0.08*x - 0.035
+ """
+
+ @staticmethod
+ def forward(ctx, x: Tensor) -> Tensor:
+ requires_grad = x.requires_grad
+ if x.dtype == torch.float16 or x.dtype == torch.bfloat16:
+ x = x.to(torch.float32)
+
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+
+ coeff = -0.08
+
+ with torch.cuda.amp.autocast(enabled=False):
+ with torch.enable_grad():
+ x = x.detach()
+ x.requires_grad = True
+ y = torch.logaddexp(zero, x - 4.0) + coeff * x - 0.035
+
+ if not requires_grad:
+ return y
+
+ y.backward(gradient=torch.ones_like(y))
+
+ grad = x.grad
+ floor = coeff
+ ceil = 1.0 + coeff + 0.005
+
+ d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
+ grad
+ )
+ if __name__ == "__main__":
+ # for self-testing only.
+ assert d_scaled.min() >= 0.0
+ assert d_scaled.max() < 256.0
+
+ d_int = d_scaled.to(torch.uint8)
+ ctx.save_for_backward(d_int)
+ if x.dtype == torch.float16 or torch.is_autocast_enabled():
+ y = y.to(torch.get_autocast_gpu_dtype())
+ return y
+
+ @staticmethod
+ def backward(ctx, y_grad: Tensor) -> Tensor:
+ (d,) = ctx.saved_tensors
+ # the same constants as used in forward pass.
+
+ coeff = -0.08
+ floor = coeff
+ ceil = 1.0 + coeff + 0.005
+ d = d * ((ceil - floor) / 255.0) + floor
+ return y_grad * d
+
+
+class SwooshL(torch.nn.Module):
+ def forward(self, x: Tensor) -> Tensor:
+ """Return Swoosh-L activation."""
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+ return logaddexp(zero, x - 4.0) - 0.08 * x - 0.035
+ if not x.requires_grad:
+ return k2.swoosh_l_forward(x)
+ else:
+ return k2.swoosh_l(x)
+ # return SwooshLFunction.apply(x)
+
+
+class SwooshLOnnx(torch.nn.Module):
+ def forward(self, x: Tensor) -> Tensor:
+ """Return Swoosh-L activation."""
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+ return logaddexp_onnx(zero, x - 4.0) - 0.08 * x - 0.035
+
+
+class SwooshRFunction(torch.autograd.Function):
+ """
+ swoosh_r(x) = log(1 + exp(x-1)) - 0.08*x - 0.313261687
+
+ derivatives are between -0.08 and 0.92.
+ """
+
+ @staticmethod
+ def forward(ctx, x: Tensor) -> Tensor:
+ requires_grad = x.requires_grad
+
+ if x.dtype == torch.float16 or x.dtype == torch.bfloat16:
+ x = x.to(torch.float32)
+
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+
+ with torch.cuda.amp.autocast(enabled=False):
+ with torch.enable_grad():
+ x = x.detach()
+ x.requires_grad = True
+ y = torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687
+
+ if not requires_grad:
+ return y
+ y.backward(gradient=torch.ones_like(y))
+
+ grad = x.grad
+ floor = -0.08
+ ceil = 0.925
+
+ d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
+ grad
+ )
+ if __name__ == "__main__":
+ # for self-testing only.
+ assert d_scaled.min() >= 0.0
+ assert d_scaled.max() < 256.0
+
+ d_int = d_scaled.to(torch.uint8)
+ ctx.save_for_backward(d_int)
+ if x.dtype == torch.float16 or torch.is_autocast_enabled():
+ y = y.to(torch.get_autocast_gpu_dtype())
+ return y
+
+ @staticmethod
+ def backward(ctx, y_grad: Tensor) -> Tensor:
+ (d,) = ctx.saved_tensors
+ # the same constants as used in forward pass.
+ floor = -0.08
+ ceil = 0.925
+ d = d * ((ceil - floor) / 255.0) + floor
+ return y_grad * d
+
+
+class SwooshR(torch.nn.Module):
+ def forward(self, x: Tensor) -> Tensor:
+ """Return Swoosh-R activation."""
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+ return logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687
+ if not x.requires_grad:
+ return k2.swoosh_r_forward(x)
+ else:
+ return k2.swoosh_r(x)
+ # return SwooshRFunction.apply(x)
+
+
+class SwooshROnnx(torch.nn.Module):
+ def forward(self, x: Tensor) -> Tensor:
+ """Return Swoosh-R activation."""
+ zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
+ return logaddexp_onnx(zero, x - 1.0) - 0.08 * x - 0.313261687
+
+
+# simple version of SwooshL that does not redefine the backprop, used in
+# ActivationDropoutAndLinearFunction.
+def SwooshLForward(x: Tensor):
+ x_offset = x - 4.0
+ log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)
+ log_sum = torch.where(log_sum == float("inf"), x_offset, log_sum)
+ return log_sum - 0.08 * x - 0.035
+
+
+# simple version of SwooshR that does not redefine the backprop, used in
+# ActivationDropoutAndLinearFunction.
+def SwooshRForward(x: Tensor):
+ x_offset = x - 1.0
+ log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)
+ log_sum = torch.where(log_sum == float("inf"), x_offset, log_sum)
+ return log_sum - 0.08 * x - 0.313261687
+
+
+class ActivationDropoutAndLinearFunction(torch.autograd.Function):
+ @staticmethod
+ @custom_fwd
+ def forward(
+ ctx,
+ x: Tensor,
+ weight: Tensor,
+ bias: Optional[Tensor],
+ activation: str,
+ dropout_p: float,
+ dropout_shared_dim: Optional[int],
+ ):
+ if dropout_p != 0.0:
+ dropout_shape = list(x.shape)
+ if dropout_shared_dim is not None:
+ dropout_shape[dropout_shared_dim] = 1
+ # else it won't be very memory efficient.
+ dropout_mask = (1.0 / (1.0 - dropout_p)) * (
+ torch.rand(*dropout_shape, device=x.device, dtype=x.dtype) > dropout_p
+ )
+ else:
+ dropout_mask = None
+
+ ctx.save_for_backward(x, weight, bias, dropout_mask)
+
+ ctx.activation = activation
+
+ forward_activation_dict = {
+ "SwooshL": k2.swoosh_l_forward,
+ "SwooshR": k2.swoosh_r_forward,
+ }
+ # it will raise a KeyError if this fails. This will be an error. We let it
+ # propagate to the user.
+ activation_func = forward_activation_dict[activation]
+ x = activation_func(x)
+ if dropout_mask is not None:
+ x = x * dropout_mask
+ x = torch.nn.functional.linear(x, weight, bias)
+ return x
+
+ @staticmethod
+ @custom_bwd
+ def backward(ctx, ans_grad: Tensor):
+ saved = ctx.saved_tensors
+ (x, weight, bias, dropout_mask) = saved
+
+ forward_and_deriv_activation_dict = {
+ "SwooshL": k2.swoosh_l_forward_and_deriv,
+ "SwooshR": k2.swoosh_r_forward_and_deriv,
+ }
+ # the following lines a KeyError if the activation is unrecognized.
+ # This will be an error. We let it propagate to the user.
+ func = forward_and_deriv_activation_dict[ctx.activation]
+
+ y, func_deriv = func(x)
+ if dropout_mask is not None:
+ y = y * dropout_mask
+ # now compute derivative of y w.r.t. weight and bias..
+ # y: (..., in_channels), ans_grad: (..., out_channels),
+ (out_channels, in_channels) = weight.shape
+
+ in_channels = y.shape[-1]
+ g = ans_grad.reshape(-1, out_channels)
+ weight_deriv = torch.matmul(g.t(), y.reshape(-1, in_channels))
+ y_deriv = torch.matmul(ans_grad, weight)
+ bias_deriv = None if bias is None else g.sum(dim=0)
+ x_deriv = y_deriv * func_deriv
+ if dropout_mask is not None:
+ # order versus func_deriv does not matter
+ x_deriv = x_deriv * dropout_mask
+
+ return x_deriv, weight_deriv, bias_deriv, None, None, None
+
+
+class ActivationDropoutAndLinear(torch.nn.Module):
+ """
+ This merges an activation function followed by dropout and then a nn.Linear module;
+ it does so in a memory efficient way so that it only stores the input to the whole
+ module. If activation == SwooshL and dropout_shared_dim != None, this will be
+ equivalent to:
+ nn.Sequential(SwooshL(),
+ Dropout3(dropout_p, shared_dim=dropout_shared_dim),
+ ScaledLinear(in_channels, out_channels, bias=bias,
+ initial_scale=initial_scale))
+ If dropout_shared_dim is None, the dropout would be equivalent to
+ Dropout2(dropout_p). Note: Dropout3 will be more memory efficient as the dropout
+ mask is smaller.
+
+ Args:
+ in_channels: number of input channels, e.g. 256
+ out_channels: number of output channels, e.g. 256
+ bias: if true, have a bias
+ activation: the activation function, for now just support SwooshL.
+ dropout_p: the dropout probability or schedule (happens after nonlinearity).
+ dropout_shared_dim: the dimension, if any, across which the dropout mask is
+ shared (e.g. the time dimension). If None, this may be less memory
+ efficient if there are modules before this one that cache the input
+ for their backprop (e.g. Balancer or Whiten).
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ bias: bool = True,
+ activation: str = "SwooshL",
+ dropout_p: FloatLike = 0.0,
+ dropout_shared_dim: Optional[int] = -1,
+ initial_scale: float = 1.0,
+ ):
+ super().__init__()
+ # create a temporary module of nn.Linear that we'll steal the
+ # weights and bias from
+ l = ScaledLinear(
+ in_channels, out_channels, bias=bias, initial_scale=initial_scale
+ )
+
+ self.weight = l.weight
+ # register_parameter properly handles making it a parameter when l.bias
+ # is None. I think there is some reason for doing it this way rather
+ # than just setting it to None but I don't know what it is, maybe
+ # something to do with exporting the module..
+ self.register_parameter("bias", l.bias)
+
+ self.activation = activation
+ self.dropout_p = dropout_p
+ self.dropout_shared_dim = dropout_shared_dim
+
+ def forward(self, x: Tensor):
+ if not self.training or torch.jit.is_scripting() or torch.jit.is_tracing():
+ if self.activation == "SwooshL":
+ x = SwooshLForward(x)
+ elif self.activation == "SwooshR":
+ x = SwooshRForward(x)
+ else:
+ assert False, self.activation
+ return torch.nn.functional.linear(x, self.weight, self.bias)
+
+ return ActivationDropoutAndLinearFunction.apply(
+ x,
+ self.weight,
+ self.bias,
+ self.activation,
+ float(self.dropout_p),
+ self.dropout_shared_dim,
+ )
+
+
+def convert_num_channels(x: Tensor, num_channels: int) -> Tensor:
+ if num_channels <= x.shape[-1]:
+ return x[..., :num_channels]
+ else:
+ shape = list(x.shape)
+ shape[-1] = num_channels - shape[-1]
+ zeros = torch.zeros(shape, dtype=x.dtype, device=x.device)
+ return torch.cat((x, zeros), dim=-1)
+
+
+def _test_whiten():
+ for proportion in [0.1, 0.5, 10.0]:
+ logging.info(f"_test_whiten(): proportion = {proportion}")
+ x = torch.randn(100, 128)
+ direction = torch.randn(128)
+ coeffs = torch.randn(100, 1)
+ x += proportion * direction * coeffs
+
+ x.requires_grad = True
+
+ m = Whiten(
+ 1, 5.0, prob=1.0, grad_scale=0.1 # num_groups # whitening_limit,
+ ) # grad_scale
+
+ for _ in range(4):
+ y = m(x)
+
+ y_grad = torch.randn_like(x)
+ y.backward(gradient=y_grad)
+
+ if proportion < 0.2:
+ assert torch.allclose(x.grad, y_grad)
+ elif proportion > 1.0:
+ assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_balancer_sign():
+ probs = torch.arange(0, 1, 0.01)
+ N = 1000
+ x = 1.0 * ((2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0)
+ x = x.detach()
+ x.requires_grad = True
+ m = Balancer(
+ probs.numel(),
+ channel_dim=0,
+ min_positive=0.05,
+ max_positive=0.95,
+ min_abs=0.0,
+ prob=1.0,
+ )
+
+ y_grad = torch.sign(torch.randn(probs.numel(), N))
+
+ y = m(x)
+ y.backward(gradient=y_grad)
+ print("_test_balancer_sign: x = ", x)
+ print("_test_balancer_sign: y grad = ", y_grad)
+ print("_test_balancer_sign: x grad = ", x.grad)
+
+
+def _test_balancer_magnitude():
+ magnitudes = torch.arange(0, 1, 0.01)
+ N = 1000
+ x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(-1)
+ x = x.detach()
+ x.requires_grad = True
+ m = Balancer(
+ magnitudes.numel(),
+ channel_dim=0,
+ min_positive=0.0,
+ max_positive=1.0,
+ min_abs=0.2,
+ max_abs=0.7,
+ prob=1.0,
+ )
+
+ y_grad = torch.sign(torch.randn(magnitudes.numel(), N))
+
+ y = m(x)
+ y.backward(gradient=y_grad)
+ print("_test_balancer_magnitude: x = ", x)
+ print("_test_balancer_magnitude: y grad = ", y_grad)
+ print("_test_balancer_magnitude: x grad = ", x.grad)
+
+
+def _test_double_swish_deriv():
+ x = torch.randn(10, 12, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ m = DoubleSwish()
+
+ tol = (1.2 - (-0.043637)) / 255.0
+ torch.autograd.gradcheck(m, x, atol=tol)
+
+ # for self-test.
+ x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ y = m(x)
+
+
+def _test_swooshl_deriv():
+ x = torch.randn(10, 12, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ m = SwooshL()
+
+ tol = 1.0 / 255.0
+ torch.autograd.gradcheck(m, x, atol=tol, eps=0.01)
+
+ # for self-test.
+ x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ y = m(x)
+
+
+def _test_swooshr_deriv():
+ x = torch.randn(10, 12, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ m = SwooshR()
+
+ tol = 1.0 / 255.0
+ torch.autograd.gradcheck(m, x, atol=tol, eps=0.01)
+
+ # for self-test.
+ x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+ x.requires_grad = True
+ y = m(x)
+
+
+def _test_softmax():
+ a = torch.randn(2, 10, dtype=torch.float64)
+ b = a.clone()
+ a.requires_grad = True
+ b.requires_grad = True
+ a.softmax(dim=1)[:, 0].sum().backward()
+ print("a grad = ", a.grad)
+ softmax(b, dim=1)[:, 0].sum().backward()
+ print("b grad = ", b.grad)
+ assert torch.allclose(a.grad, b.grad)
+
+
+def _test_piecewise_linear():
+ p = PiecewiseLinear((0, 10.0))
+ for x in [-100, 0, 100]:
+ assert p(x) == 10.0
+ p = PiecewiseLinear((0, 10.0), (1, 0.0))
+ for x, y in [(-100, 10.0), (0, 10.0), (0.5, 5.0), (1, 0.0), (2, 0.0)]:
+ print("x, y = ", x, y)
+ assert p(x) == y, (x, p(x), y)
+
+ q = PiecewiseLinear((0.5, 15.0), (0.6, 1.0))
+ x_vals = [-1.0, 0.0, 0.1, 0.2, 0.5, 0.6, 0.7, 0.9, 1.0, 2.0]
+ pq = p.max(q)
+ for x in x_vals:
+ y1 = max(p(x), q(x))
+ y2 = pq(x)
+ assert abs(y1 - y2) < 0.001
+ pq = p.min(q)
+ for x in x_vals:
+ y1 = min(p(x), q(x))
+ y2 = pq(x)
+ assert abs(y1 - y2) < 0.001
+ pq = p + q
+ for x in x_vals:
+ y1 = p(x) + q(x)
+ y2 = pq(x)
+ assert abs(y1 - y2) < 0.001
+
+
+def _test_activation_dropout_and_linear():
+ in_channels = 20
+ out_channels = 30
+
+ for bias in [True, False]:
+ # actually we don't test for dropout_p != 0.0 because forward functions will give
+ # different answers. This is because we are using the k2 implementation of
+ # swoosh_l an swoosh_r inside SwooshL() and SwooshR(), and they call randn()
+ # internally, messing up the random state.
+ for dropout_p in [0.0]:
+ for activation in ["SwooshL", "SwooshR"]:
+ m1 = nn.Sequential(
+ SwooshL() if activation == "SwooshL" else SwooshR(),
+ Dropout3(p=dropout_p, shared_dim=-1),
+ ScaledLinear(
+ in_channels, out_channels, bias=bias, initial_scale=0.5
+ ),
+ )
+ m2 = ActivationDropoutAndLinear(
+ in_channels,
+ out_channels,
+ bias=bias,
+ initial_scale=0.5,
+ activation=activation,
+ dropout_p=dropout_p,
+ )
+ with torch.no_grad():
+ m2.weight[:] = m1[2].weight
+ if bias:
+ m2.bias[:] = m1[2].bias
+ # make sure forward gives same result.
+ x1 = torch.randn(10, in_channels)
+ x1.requires_grad = True
+
+ # TEMP.
+ assert torch.allclose(
+ SwooshRFunction.apply(x1), SwooshRForward(x1), atol=1.0e-03
+ )
+
+ x2 = x1.clone().detach()
+ x2.requires_grad = True
+ seed = 10
+ torch.manual_seed(seed)
+ y1 = m1(x1)
+ y_grad = torch.randn_like(y1)
+ y1.backward(gradient=y_grad)
+ torch.manual_seed(seed)
+ y2 = m2(x2)
+ y2.backward(gradient=y_grad)
+
+ print(
+ f"bias = {bias}, dropout_p = {dropout_p}, activation = {activation}"
+ )
+ print("y1 = ", y1)
+ print("y2 = ", y2)
+ assert torch.allclose(y1, y2, atol=0.02)
+ assert torch.allclose(m1[2].weight.grad, m2.weight.grad, atol=1.0e-05)
+ if bias:
+ assert torch.allclose(m1[2].bias.grad, m2.bias.grad, atol=1.0e-05)
+ print("x1.grad = ", x1.grad)
+ print("x2.grad = ", x2.grad)
+
+ def isclose(a, b):
+ # return true if cosine similarity is > 0.9.
+ return (a * b).sum() > 0.9 * (
+ (a**2).sum() * (b**2).sum()
+ ).sqrt()
+
+ # the SwooshL() implementation has a noisy gradient due to 1-byte
+ # storage of it.
+ assert isclose(x1.grad, x2.grad)
+
+
+if __name__ == "__main__":
+ logging.getLogger().setLevel(logging.INFO)
+ torch.set_num_threads(1)
+ torch.set_num_interop_threads(1)
+ _test_piecewise_linear()
+ _test_softmax()
+ _test_whiten()
+ _test_balancer_sign()
+ _test_balancer_magnitude()
+ _test_double_swish_deriv()
+ _test_swooshr_deriv()
+ _test_swooshl_deriv()
+ _test_activation_dropout_and_linear()
diff --git a/egs/europarl_st/SRT/lcma_srt/subsampling.py b/egs/europarl_st/SRT/lcma_srt/subsampling.py
new file mode 100644
index 0000000000..b2f769d3f6
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/subsampling.py
@@ -0,0 +1,406 @@
+#!/usr/bin/env python3
+# Copyright 2023 Xiaomi Corp. (authors: Daniel Povey,
+# Zengwei Yao)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+from typing import Tuple
+
+import torch
+from scaling import (
+ Balancer,
+ BiasNorm,
+ Dropout3,
+ FloatLike,
+ Optional,
+ ScaledConv2d,
+ ScaleGrad,
+ ScheduledFloat,
+ SwooshL,
+ SwooshR,
+ Whiten,
+)
+from torch import Tensor, nn
+
+
+class ConvNeXt(nn.Module):
+ """
+ Our interpretation of the ConvNeXt module as used in https://arxiv.org/pdf/2206.14747.pdf
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ hidden_ratio: int = 3,
+ kernel_size: Tuple[int, int] = (7, 7),
+ layerdrop_rate: FloatLike = None,
+ ):
+ super().__init__()
+ self.padding = ((kernel_size[0] - 1) // 2, (kernel_size[1] - 1) // 2)
+ hidden_channels = channels * hidden_ratio
+ if layerdrop_rate is None:
+ layerdrop_rate = ScheduledFloat((0.0, 0.2), (20000.0, 0.015))
+ self.layerdrop_rate = layerdrop_rate
+
+ self.depthwise_conv = nn.Conv2d(
+ in_channels=channels,
+ out_channels=channels,
+ groups=channels,
+ kernel_size=kernel_size,
+ padding=self.padding,
+ )
+
+ self.pointwise_conv1 = nn.Conv2d(
+ in_channels=channels, out_channels=hidden_channels, kernel_size=1
+ )
+
+ self.hidden_balancer = Balancer(
+ hidden_channels,
+ channel_dim=1,
+ min_positive=0.3,
+ max_positive=1.0,
+ min_abs=0.75,
+ max_abs=5.0,
+ )
+
+ self.activation = SwooshL()
+ self.pointwise_conv2 = ScaledConv2d(
+ in_channels=hidden_channels,
+ out_channels=channels,
+ kernel_size=1,
+ initial_scale=0.01,
+ )
+
+ self.out_balancer = Balancer(
+ channels,
+ channel_dim=1,
+ min_positive=0.4,
+ max_positive=0.6,
+ min_abs=1.0,
+ max_abs=6.0,
+ )
+ self.out_whiten = Whiten(
+ num_groups=1,
+ whitening_limit=5.0,
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ def forward(self, x: Tensor) -> Tensor:
+ if torch.jit.is_scripting() or torch.jit.is_tracing() or not self.training:
+ return self.forward_internal(x)
+ layerdrop_rate = float(self.layerdrop_rate)
+
+ if layerdrop_rate != 0.0:
+ batch_size = x.shape[0]
+ mask = (
+ torch.rand((batch_size, 1, 1, 1), dtype=x.dtype, device=x.device)
+ > layerdrop_rate
+ )
+ else:
+ mask = None
+ # turns out this caching idea does not work with --world-size > 1
+ # return caching_eval(self.forward_internal, x, mask)
+ return self.forward_internal(x, mask)
+
+ def forward_internal(
+ self, x: Tensor, layer_skip_mask: Optional[Tensor] = None
+ ) -> Tensor:
+ """
+ x layout: (N, C, H, W), i.e. (batch_size, num_channels, num_frames, num_freqs)
+
+ The returned value has the same shape as x.
+ """
+ bypass = x
+ x = self.depthwise_conv(x)
+ x = self.pointwise_conv1(x)
+ x = self.hidden_balancer(x)
+ x = self.activation(x)
+ x = self.pointwise_conv2(x)
+
+ if layer_skip_mask is not None:
+ x = x * layer_skip_mask
+
+ x = bypass + x
+ x = self.out_balancer(x)
+
+ if x.requires_grad:
+ x = x.transpose(1, 3) # (N, W, H, C); need channel dim to be last
+ x = self.out_whiten(x)
+ x = x.transpose(1, 3) # (N, C, H, W)
+
+ return x
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ cached_left_pad: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ """
+ Args:
+ x layout: (N, C, H, W), i.e. (batch_size, num_channels, num_frames, num_freqs)
+ cached_left_pad: (batch_size, num_channels, left_pad, num_freqs)
+
+ Returns:
+ - The returned value has the same shape as x.
+ - Updated cached_left_pad.
+ """
+ padding = self.padding
+
+ # The length without right padding for depth-wise conv
+ T = x.size(2) - padding[0]
+
+ bypass = x[:, :, :T, :]
+
+ # Pad left side
+ assert cached_left_pad.size(2) == padding[0], (
+ cached_left_pad.size(2),
+ padding[0],
+ )
+ x = torch.cat([cached_left_pad, x], dim=2)
+ # Update cached left padding
+ cached_left_pad = x[:, :, T : padding[0] + T, :]
+
+ # depthwise_conv
+ x = torch.nn.functional.conv2d(
+ x,
+ weight=self.depthwise_conv.weight,
+ bias=self.depthwise_conv.bias,
+ padding=(0, padding[1]),
+ groups=self.depthwise_conv.groups,
+ )
+ x = self.pointwise_conv1(x)
+ x = self.hidden_balancer(x)
+ x = self.activation(x)
+ x = self.pointwise_conv2(x)
+
+ x = bypass + x
+ return x, cached_left_pad
+
+
+class Conv2dSubsampling(nn.Module):
+ """Convolutional 2D subsampling (to 1/2 length).
+
+ Convert an input of shape (N, T, idim) to an output
+ with shape (N, T', odim), where
+ T' = (T-3)//2 - 2 == (T-7)//2
+
+ It is based on
+ https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/transformer/subsampling.py # noqa
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ layer1_channels: int = 8,
+ layer2_channels: int = 32,
+ layer3_channels: int = 128,
+ dropout: FloatLike = 0.1,
+ ) -> None:
+ """
+ Args:
+ in_channels:
+ Number of channels in. The input shape is (N, T, in_channels).
+ Caution: It requires: T >=7, in_channels >=7
+ out_channels
+ Output dim. The output shape is (N, (T-3)//2, out_channels)
+ layer1_channels:
+ Number of channels in layer1
+ layer1_channels:
+ Number of channels in layer2
+ bottleneck:
+ bottleneck dimension for 1d squeeze-excite
+ """
+ assert in_channels >= 7
+ super().__init__()
+
+ # The ScaleGrad module is there to prevent the gradients
+ # w.r.t. the weight or bias of the first Conv2d module in self.conv from
+ # exceeding the range of fp16 when using automatic mixed precision (amp)
+ # training. (The second one is necessary to stop its bias from getting
+ # a too-large gradient).
+
+ self.conv = nn.Sequential(
+ nn.Conv2d(
+ in_channels=1,
+ out_channels=layer1_channels,
+ kernel_size=3,
+ padding=(0, 1), # (time, freq)
+ ),
+ ScaleGrad(0.2),
+ Balancer(layer1_channels, channel_dim=1, max_abs=1.0),
+ SwooshR(),
+ nn.Conv2d(
+ in_channels=layer1_channels,
+ out_channels=layer2_channels,
+ kernel_size=3,
+ stride=2,
+ padding=0,
+ ),
+ Balancer(layer2_channels, channel_dim=1, max_abs=4.0),
+ SwooshR(),
+ nn.Conv2d(
+ in_channels=layer2_channels,
+ out_channels=layer3_channels,
+ kernel_size=3,
+ stride=(1, 2), # (time, freq)
+ ),
+ Balancer(layer3_channels, channel_dim=1, max_abs=4.0),
+ SwooshR(),
+ )
+
+ # just one convnext layer
+ self.convnext = ConvNeXt(layer3_channels, kernel_size=(7, 7))
+
+ # (in_channels-3)//4
+ self.out_width = (((in_channels - 1) // 2) - 1) // 2
+ self.layer3_channels = layer3_channels
+
+ self.out = nn.Linear(self.out_width * layer3_channels, out_channels)
+ # use a larger than normal grad_scale on this whitening module; there is
+ # only one such module, so there is not a concern about adding together
+ # many copies of this extra gradient term.
+ self.out_whiten = Whiten(
+ num_groups=1,
+ whitening_limit=ScheduledFloat((0.0, 4.0), (20000.0, 8.0), default=4.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.02,
+ )
+
+ # max_log_eps=0.0 is to prevent both eps and the output of self.out from
+ # getting large, there is an unnecessary degree of freedom.
+ self.out_norm = BiasNorm(out_channels)
+ self.dropout = Dropout3(dropout, shared_dim=1)
+
+ def forward(
+ self, x: torch.Tensor, x_lens: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Subsample x.
+
+ Args:
+ x:
+ Its shape is (N, T, idim).
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+
+ Returns:
+ - a tensor of shape (N, (T-7)//2, odim)
+ - output lengths, of shape (batch_size,)
+ """
+ # On entry, x is (N, T, idim)
+ x = x.unsqueeze(1) # (N, T, idim) -> (N, 1, T, idim) i.e., (N, C, H, W)
+ # scaling x by 0.1 allows us to use a larger grad-scale in fp16 "amp" (automatic mixed precision)
+ # training, since the weights in the first convolution are otherwise the limiting factor for getting infinite
+ # gradients.
+ x = self.conv(x)
+ x = self.convnext(x)
+
+ # Now x is of shape (N, odim, (T-7)//2, (idim-3)//4)
+ b, c, t, f = x.size()
+
+ x = x.transpose(1, 2).reshape(b, t, c * f)
+ # now x: (N, (T-7)//2, out_width * layer3_channels))
+
+ x = self.out(x)
+ # Now x is of shape (N, (T-7)//2, odim)
+ x = self.out_whiten(x)
+ x = self.out_norm(x)
+ x = self.dropout(x)
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ x_lens = (x_lens - 7) // 2
+ else:
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ x_lens = (x_lens - 7) // 2
+ assert x.size(1) == x_lens.max().item(), (x.size(1), x_lens.max())
+
+ return x, x_lens
+
+ def streaming_forward(
+ self,
+ x: torch.Tensor,
+ x_lens: torch.Tensor,
+ cached_left_pad: Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """Subsample x.
+
+ Args:
+ x:
+ Its shape is (N, T, idim).
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+
+ Returns:
+ - a tensor of shape (N, (T-7)//2, odim)
+ - output lengths, of shape (batch_size,)
+ - updated cache
+ """
+ # On entry, x is (N, T, idim)
+ x = x.unsqueeze(1) # (N, T, idim) -> (N, 1, T, idim) i.e., (N, C, H, W)
+
+ # T' = (T-7)//2
+ x = self.conv(x)
+
+ # T' = (T-7)//2-3
+ x, cached_left_pad = self.convnext.streaming_forward(
+ x, cached_left_pad=cached_left_pad
+ )
+
+ # Now x is of shape (N, odim, T', ((idim-1)//2 - 1)//2)
+ b, c, t, f = x.size()
+
+ x = x.transpose(1, 2).reshape(b, t, c * f)
+ # now x: (N, T', out_width * layer3_channels))
+
+ x = self.out(x)
+ # Now x is of shape (N, T', odim)
+ x = self.out_norm(x)
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ assert self.convnext.padding[0] == 3
+ # The ConvNeXt module needs 3 frames of right padding after subsampling
+ x_lens = (x_lens - 7) // 2 - 3
+ else:
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ # The ConvNeXt module needs 3 frames of right padding after subsampling
+ assert self.convnext.padding[0] == 3
+ x_lens = (x_lens - 7) // 2 - 3
+
+ assert x.size(1) == x_lens.max().item(), (x.shape, x_lens.max())
+
+ return x, x_lens, cached_left_pad
+
+ @torch.jit.export
+ def get_init_states(
+ self,
+ batch_size: int = 1,
+ device: torch.device = torch.device("cpu"),
+ ) -> Tensor:
+ """Get initial states for Conv2dSubsampling module.
+ It is the cached left padding for ConvNeXt module,
+ of shape (batch_size, num_channels, left_pad, num_freqs)
+ """
+ left_pad = self.convnext.padding[0]
+ freq = self.out_width
+ channels = self.layer3_channels
+ cached_embed_left_pad = torch.zeros(batch_size, channels, left_pad, freq).to(
+ device
+ )
+
+ return cached_embed_left_pad
diff --git a/egs/europarl_st/SRT/lcma_srt/train.py b/egs/europarl_st/SRT/lcma_srt/train.py
new file mode 100644
index 0000000000..80ed7d80cf
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train.py
@@ -0,0 +1,2196 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+import argparse
+import copy
+import logging
+import os
+import random
+import re
+import sys
+import time
+import warnings
+from collections import defaultdict
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import k2
+import numpy as np
+import sentencepiece as spm
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+import torch.nn as nn
+from attention_decoder import AttentionDecoderModel
+
+# Local project modules
+from datamodule import LibriSpeechAsrDataModule
+from decoder import Decoder
+from encoder_interface import EncoderInterface
+from joiner import Joiner
+from lhotse import CutSet, load_manifest_lazy
+
+# Data utils
+from lhotse.cut import Cut
+from lhotse.dataset import SpecAugment
+from lhotse.dataset.sampling.base import CutSampler
+from lhotse.utils import fix_random_seed
+
+# LCMA-SRT model
+from model import LCMASRTModel
+from optim import Eden, LRScheduler, ScaledAdam
+from scaling import ScheduledFloat
+from subsampling import Conv2dSubsampling
+from torch import Tensor
+from torch.cuda.amp import GradScaler
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.tensorboard import SummaryWriter
+from zipformer import Zipformer2
+
+# Icefall / project deps
+from icefall import diagnostics
+from icefall.checkpoint import load_checkpoint, remove_checkpoints
+from icefall.checkpoint import save_checkpoint as save_checkpoint_impl
+from icefall.checkpoint import (
+ save_checkpoint_with_global_batch_idx,
+ update_averaged_model,
+)
+from icefall.dist import cleanup_dist, setup_dist
+from icefall.env import get_env_info
+from icefall.err import raise_grad_scale_is_too_small_error
+from icefall.hooks import register_inf_check_hooks
+from icefall.utils import (
+ AttributeDict,
+ MetricsTracker,
+ get_parameter_groups_with_lrs,
+ setup_logger,
+ str2bool,
+)
+
+LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, LRScheduler]
+
+
+# ------------------------------
+# Helpers
+# ------------------------------
+
+
+def _to_int_tuple(s: str) -> Tuple[int, ...]:
+ return tuple(map(int, s.split(",")))
+
+
+def load_cuts_lazy(manifest_paths: str, shuffle: bool = False) -> CutSet:
+ paths = [path.strip() for path in re.split(r"[,;]", manifest_paths) if path.strip()]
+ if not paths:
+ raise ValueError("No valid manifest paths provided.")
+
+ logging.info(f"Loading CutSets lazily from {len(paths)} manifest files.")
+
+ combined_cuts = load_manifest_lazy(paths[0])
+ for path in paths[1:]:
+ cuts = load_manifest_lazy(path)
+ combined_cuts = combined_cuts + cuts
+ # Optionally shuffle the combined CutSet
+ if shuffle:
+ logging.info(f"## start to shuffle the feature ...")
+ start_time = time.time()
+ combined_cuts = combined_cuts.shuffle()
+ elapsed = time.time() - start_time
+ logging.info(f"Shuffling took {elapsed:.2f} seconds.")
+
+ return combined_cuts
+
+
+def get_adjusted_batch_count(params: AttributeDict) -> float:
+ return (
+ params.batch_idx_train
+ * (params.max_duration * params.world_size)
+ / params.ref_duration
+ )
+
+
+def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None:
+ if isinstance(model, DDP):
+ model = model.module
+ for name, module in model.named_modules():
+ if hasattr(module, "batch_count"):
+ module.batch_count = batch_count
+ if hasattr(module, "name"):
+ module.name = name
+
+
+def _capture_rng_state() -> Dict[str, Any]:
+ state: Dict[str, Any] = {
+ "python": random.getstate(),
+ "torch": torch.random.get_rng_state(),
+ }
+ try:
+ state["numpy"] = np.random.get_state()
+ except Exception:
+ pass
+ if torch.cuda.is_available():
+ state["cuda"] = torch.cuda.get_rng_state_all()
+ return state
+
+
+def _restore_rng_state(state: Optional[Dict[str, Any]]) -> None:
+ if not state:
+ return
+ if "python" in state:
+ random.setstate(state["python"])
+ if "numpy" in state:
+ try:
+ np.random.set_state(state["numpy"])
+ except Exception:
+ logging.warning("Failed to restore NumPy RNG state")
+ if "torch" in state:
+ torch.random.set_rng_state(state["torch"])
+ if "cuda" in state and torch.cuda.is_available():
+ torch.cuda.set_rng_state_all(state["cuda"])
+
+
+_RESUME_STATE_KEYS = [
+ "batch_idx_train",
+ "cur_epoch",
+ "best_train_epoch",
+ "best_train_loss",
+ "best_valid_epoch",
+ "best_valid_loss",
+ "start_epoch",
+ "start_batch",
+]
+
+
+def _collect_resume_state(params: AttributeDict) -> Dict[str, Any]:
+ state: Dict[str, Any] = {}
+ for key in _RESUME_STATE_KEYS:
+ if hasattr(params, key):
+ state[key] = getattr(params, key)
+ if hasattr(params, "rng_state"):
+ state["rng_state"] = getattr(params, "rng_state")
+ return state
+
+
+def _build_params_payload(params: Optional[AttributeDict]) -> Optional[Dict[str, Any]]:
+ if params is None:
+ return None
+ if hasattr(params, "items"):
+ params_to_save = dict(params.items())
+ else:
+ params_to_save = dict(params)
+ params_to_save["params"] = _collect_resume_state(params)
+ return params_to_save
+
+
+# ------------------------------
+# Argparse & default params
+# ------------------------------
+
+
+def add_encoder_args(group: argparse._ArgumentGroup, prefix: str = "asr") -> None:
+ """Add Zipformer encoder args for either ASR or ST with a prefix."""
+ pf = prefix
+ group.add_argument(
+ f"--num-encoder-layers-{pf}",
+ type=str,
+ default="2,2,3,4,3,2",
+ help="Number of zipformer encoder layers per stack, comma separated.",
+ )
+ group.add_argument(
+ f"--downsampling-factor-{pf}",
+ type=str,
+ default="1,2,4,8,4,2",
+ help="Downsampling factor for each stack of encoder layers.",
+ )
+ group.add_argument(
+ f"--feedforward-dim-{pf}",
+ type=str,
+ default="512,768,1024,1536,1024,768",
+ help="Feedforward dimension per stack, comma separated.",
+ )
+ group.add_argument(
+ f"--num-heads-{pf}",
+ type=str,
+ default="4,4,4,8,4,4",
+ help="Number of attention heads: a single int or comma-separated list.",
+ )
+ group.add_argument(
+ f"--encoder-dim-{pf}",
+ type=str,
+ default="192,256,384,512,384,256",
+ help="Embedding dimension per stack: int or comma-separated list.",
+ )
+ group.add_argument(
+ f"--query-head-dim-{pf}",
+ type=str,
+ default="32",
+ help="Query/key dim per head: int or comma-separated list.",
+ )
+ group.add_argument(
+ f"--value-head-dim-{pf}",
+ type=str,
+ default="12",
+ help="Value dim per head: int or comma-separated list.",
+ )
+ group.add_argument(
+ f"--pos-head-dim-{pf}",
+ type=str,
+ default="4",
+ help="Positional-encoding dim per head: int or comma-separated list.",
+ )
+ group.add_argument(
+ f"--pos-dim-{pf}",
+ type=int,
+ default=48,
+ help="Positional-encoding embedding dimension",
+ )
+ group.add_argument(
+ f"--encoder-unmasked-dim-{pf}",
+ type=str,
+ default="192,192,256,256,256,192",
+ help=(
+ "Unmasked dims in encoders for dropout aug; int or CSV; must be "
+ "<= corresponding encoder_dim."
+ ),
+ )
+ group.add_argument(
+ f"--cnn-module-kernel-{pf}",
+ type=str,
+ default="31,31,15,15,15,31",
+ help="Kernel sizes in conv modules per stack: int or CSV.",
+ )
+
+
+def add_model_arguments(parser: argparse.ArgumentParser) -> None:
+ # Freeze ASR
+ parser.add_argument("--freeze-asr", type=str2bool, default=False)
+ parser.add_argument("--freeze-frontend", type=str2bool, default=False)
+
+ # Shared frontend
+ parser.add_argument(
+ "--feature-dim",
+ type=int,
+ default=80,
+ help="Input feature dim (must match FBank config).",
+ )
+ parser.add_argument("--subsampling-factor", type=int, default=4)
+
+ # ASR encoder & head
+ add_encoder_args(parser.add_argument_group("ASR encoder"), prefix="asr")
+ parser.add_argument("--decoder-dim-asr", type=int, default=512)
+ parser.add_argument("--joiner-dim-asr", type=int, default=512)
+ parser.add_argument("--output-downsampling-factor-asr", type=int, default=2)
+
+ # ST encoder & head (stacked on top of ASR encoder)
+ add_encoder_args(parser.add_argument_group("ST encoder"), prefix="st")
+ parser.add_argument("--decoder-dim-st", type=int, default=512)
+ parser.add_argument("--joiner-dim-st", type=int, default=512)
+ parser.add_argument("--output-downsampling-factor-st", type=int, default=2)
+
+ # Attention decoders (optional)
+ parser.add_argument("--use-attention-decoder", type=str2bool, default=False)
+ parser.add_argument("--attention-decoder-dim-asr", type=int, default=512)
+ parser.add_argument("--attention-decoder-dim-st", type=int, default=512)
+ parser.add_argument("--attention-decoder-num-layers", type=int, default=6)
+ parser.add_argument("--attention-decoder-attention-dim", type=int, default=512)
+ parser.add_argument("--attention-decoder-num-heads", type=int, default=8)
+ parser.add_argument("--attention-decoder-feedforward-dim", type=int, default=2048)
+
+ # Causal / streaming
+ parser.add_argument("--causal", type=str2bool, default=False)
+ parser.add_argument(
+ "--chunk-size",
+ type=str,
+ default="16,32,64,-1",
+ help="Chunk sizes (at 50Hz) for streaming; must be just -1 if non-causal.",
+ )
+ parser.add_argument(
+ "--left-context-frames",
+ type=str,
+ default="64,128,256,-1",
+ help="Max left context in frames for causal training.",
+ )
+
+ # Heads toggles & losses
+ parser.add_argument("--use-transducer", type=str2bool, default=True)
+ parser.add_argument("--use-ctc-asr", type=str2bool, default=False)
+ parser.add_argument("--use-ctc-st", type=str2bool, default=False)
+ parser.add_argument("--use-cr-ctc", type=str2bool, default=False)
+ parser.add_argument("--lm-scale", type=float, default=0.25)
+ parser.add_argument("--am-scale", type=float, default=0.0)
+ parser.add_argument("--simple-loss-scale", type=float, default=0.5)
+ parser.add_argument("--ctc-loss-scale", type=float, default=0.2)
+ parser.add_argument("--cr-loss-scale", type=float, default=0.2)
+ parser.add_argument("--time-mask-ratio", type=float, default=2.5)
+ parser.add_argument("--attention-decoder-loss-scale", type=float, default=0.8)
+
+ # Task weights & distillation
+ parser.add_argument("--task-weight-asr", type=float, default=1.0)
+ parser.add_argument("--task-weight-st", type=float, default=1.0)
+ parser.add_argument(
+ "--self-distill-scale",
+ type=float,
+ default=0.0,
+ help="If the model returns a distillation loss, scale it by this factor.",
+ )
+
+ # Tokenizers
+ parser.add_argument(
+ "--bpe-model-asr", type=str, default="data/lang_bpe_500/bpe.model"
+ )
+ parser.add_argument(
+ "--bpe-model-st", type=str, default="data/lang_st_bpe_1k/bpe.model"
+ )
+ # parser.add_argument("--ast-use-asr-data", type=int, default=0)
+ # parser.add_argument("--asr-use-ast-data", type=int, default=0)
+
+ parser.add_argument(
+ "--tgt-langs",
+ type=str,
+ default="en,zh-cn,es",
+ help="Comma-separated target language tags aligned with <2xx> labels, e.g. zh-cn,ja,de",
+ )
+
+ parser.add_argument(
+ "--srt-langs",
+ type=str,
+ default="en,zh-cn",
+ help="Comma-separated target language tags aligned with <2xx> labels, e.g. zh-cn,ja,de",
+ )
+
+ parser.add_argument(
+ "--enable-st",
+ type=str2bool,
+ default=True,
+ help="Whether to enable ST branch (encoder + losses). "
+ "For 1st stage ASR-only pretraining, set to 0.",
+ )
+
+ parser.add_argument(
+ "--asr-moe", type=str2bool, default=True, help="ASR: enable MoE adapter."
+ )
+ parser.add_argument(
+ "--asr-src",
+ type=str2bool,
+ default=True,
+ help="ASR: use source language ID (in MoE routing if asr-moe, otherwise bias).",
+ )
+ parser.add_argument(
+ "--ast-moe", type=str2bool, default=True, help="AST: enable MoE adapter."
+ )
+ parser.add_argument(
+ "--ast-tgt",
+ type=str2bool,
+ default=True,
+ help="AST: use target language ID (in MoE routing if ast-moe, otherwise bias).",
+ )
+ parser.add_argument("--num-experts-asr", type=int, default=4)
+ parser.add_argument("--num-experts-ast", type=int, default=8)
+ parser.add_argument("--entropy-reg-asr", type=float, default=0.0)
+ parser.add_argument("--entropy-reg-ast", type=float, default=0.0)
+ parser.add_argument("--temperature-asr", type=float, default=1.0)
+ parser.add_argument("--temperature-ast", type=float, default=1.0)
+
+
+def get_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
+ )
+
+ # DDP & bookkeeping
+ parser.add_argument("--world-size", type=int, default=1)
+ parser.add_argument("--master-port", type=int, default=12354)
+ parser.add_argument("--tensorboard", type=str2bool, default=True)
+
+ # Training schedule
+ parser.add_argument("--num-epochs", type=int, default=30)
+ parser.add_argument("--start-epoch", type=int, default=1)
+ parser.add_argument("--start-batch", type=int, default=0)
+ parser.add_argument(
+ "--reset-progress-stats",
+ type=str2bool,
+ default=False,
+ help="Set to 1 when loading weights for a new training stage; "
+ "skips restoring batch/epoch counters, RNG, and sampler state.",
+ )
+
+ # Paths
+ parser.add_argument("--exp-dir", type=str, default="exp")
+ parser.add_argument("--baige-tb-dir", type=str, default="exp")
+
+ # Optim & LR
+ parser.add_argument("--base-lr", type=float, default=0.045)
+ parser.add_argument("--lr-batches", type=float, default=7500)
+ parser.add_argument("--lr-epochs", type=float, default=3.5)
+ parser.add_argument("--ref-duration", type=float, default=600.0)
+
+ # Decoder context
+ parser.add_argument(
+ "--context-size",
+ type=int,
+ default=2,
+ help="RNN-T decoder context size: 1=bigram; 2=trigram.",
+ )
+
+ # RNNT pruning
+ parser.add_argument("--prune-range-asr", type=int, default=5)
+ parser.add_argument("--prune-range-st", type=int, default=10)
+
+ # Logging & misc
+ parser.add_argument("--seed", type=int, default=42)
+ parser.add_argument("--print-diagnostics", type=str2bool, default=False)
+ parser.add_argument("--skip-sanity-check", type=str2bool, default=False)
+ parser.add_argument("--inf-check", type=str2bool, default=False)
+ parser.add_argument(
+ "--dump-moe-routing-stats",
+ type=str2bool,
+ default=False,
+ )
+
+ # Checkpointing
+ parser.add_argument("--save-every-n", type=int, default=4000)
+ parser.add_argument("--keep-last-k", type=int, default=30)
+ parser.add_argument("--average-period", type=int, default=200)
+ parser.add_argument("--resume-from-checkpoint", type=str, default=None)
+ parser.add_argument("--remove-st-head", type=str2bool, default=False)
+ parser.add_argument(
+ "--resume-optimizer-scheduler-scaler", type=str2bool, default=True
+ )
+
+ # AMP
+ parser.add_argument("--use-fp16", type=str2bool, default=False)
+ parser.add_argument("--use-bf16", type=str2bool, default=False)
+
+ # Lhotse datasets
+ parser.add_argument("--train-cuts-paths", type=str, default=None)
+ parser.add_argument("--valid-cuts-paths", type=str, default=None)
+ parser.add_argument("--utterance-min-duration", type=float, default=0.3)
+ parser.add_argument("--utterance-max-duration", type=float, default=20.0)
+
+ parser.add_argument("--use-tgt", type=str2bool, default=False)
+ # SpecAug time-warp used inside model (forward)
+ # parser.add_argument("--spec-aug-time-warp-factor", type=int, default=0)
+ parser.add_argument("--warm-step", type=int, default=5000)
+ # Expose model args
+ add_model_arguments(parser)
+
+ # DataModule args
+ LibriSpeechAsrDataModule.add_arguments(parser)
+ return parser
+
+
+def get_params() -> AttributeDict:
+ params = AttributeDict(
+ {
+ "best_train_loss": float("inf"),
+ "best_valid_loss": float("inf"),
+ "best_train_epoch": -1,
+ "best_valid_epoch": -1,
+ "batch_idx_train": 0,
+ "log_interval": 50,
+ "reset_interval": 200,
+ "valid_interval": 3000,
+ "feature_dim": 80,
+ "subsampling_factor": 4,
+ "ignore_id": -1,
+ "label_smoothing": 0.1,
+ "env_info": get_env_info(),
+ }
+ )
+ return params
+
+
+# ------------------------------
+# Model builders (ASR/ST)
+# ------------------------------
+
+
+def get_encoder_embed(params: AttributeDict) -> nn.Module:
+ return Conv2dSubsampling(
+ in_channels=params.feature_dim,
+ out_channels=_to_int_tuple(params.encoder_dim_asr)[0],
+ dropout=ScheduledFloat((0.0, 0.3), (20000.0, 0.1)),
+ )
+
+
+def _build_zipformer(
+ prefix: str, output_downsampling_factor, params: AttributeDict
+) -> EncoderInterface:
+ enc = Zipformer2(
+ output_downsampling_factor=output_downsampling_factor,
+ downsampling_factor=_to_int_tuple(
+ getattr(params, f"downsampling_factor_{prefix}")
+ ),
+ num_encoder_layers=_to_int_tuple(
+ getattr(params, f"num_encoder_layers_{prefix}")
+ ),
+ encoder_dim=_to_int_tuple(getattr(params, f"encoder_dim_{prefix}")),
+ encoder_unmasked_dim=_to_int_tuple(
+ getattr(params, f"encoder_unmasked_dim_{prefix}")
+ ),
+ query_head_dim=_to_int_tuple(getattr(params, f"query_head_dim_{prefix}")),
+ pos_head_dim=_to_int_tuple(getattr(params, f"pos_head_dim_{prefix}")),
+ value_head_dim=_to_int_tuple(getattr(params, f"value_head_dim_{prefix}")),
+ pos_dim=getattr(params, f"pos_dim_{prefix}"),
+ num_heads=_to_int_tuple(getattr(params, f"num_heads_{prefix}")),
+ feedforward_dim=_to_int_tuple(getattr(params, f"feedforward_dim_{prefix}")),
+ cnn_module_kernel=_to_int_tuple(getattr(params, f"cnn_module_kernel_{prefix}")),
+ dropout=ScheduledFloat((0.0, 0.3), (20000.0, 0.1)),
+ warmup_batches=4000.0,
+ causal=params.causal,
+ chunk_size=_to_int_tuple(params.chunk_size),
+ left_context_frames=_to_int_tuple(params.left_context_frames),
+ )
+ return enc
+
+
+def get_decoder(
+ prefix: str, params: AttributeDict, vocab_size: int, blank_id: int
+) -> nn.Module:
+ return Decoder(
+ vocab_size=vocab_size,
+ decoder_dim=getattr(params, f"decoder_dim_{prefix}"),
+ blank_id=blank_id,
+ context_size=params.context_size,
+ )
+
+
+def get_joiner(prefix: str, params: AttributeDict, vocab_size: int) -> nn.Module:
+ return Joiner(
+ encoder_dim=max(_to_int_tuple(getattr(params, f"encoder_dim_{prefix}"))),
+ decoder_dim=getattr(params, f"decoder_dim_{prefix}"),
+ joiner_dim=getattr(params, f"joiner_dim_{prefix}"),
+ vocab_size=vocab_size,
+ )
+
+
+def get_attention_decoder(
+ prefix: str, params: AttributeDict, vocab_size: int, sos_id: int, eos_id: int
+) -> nn.Module:
+ return AttentionDecoderModel(
+ vocab_size=vocab_size,
+ decoder_dim=getattr(params, f"attention_decoder_dim_{prefix}"),
+ num_decoder_layers=params.attention_decoder_num_layers,
+ attention_dim=params.attention_decoder_attention_dim,
+ num_heads=params.attention_decoder_num_heads,
+ feedforward_dim=params.attention_decoder_feedforward_dim,
+ memory_dim=max(_to_int_tuple(getattr(params, f"encoder_dim_{prefix}"))),
+ sos_id=sos_id,
+ eos_id=eos_id,
+ ignore_id=params.ignore_id,
+ label_smoothing=params.label_smoothing,
+ )
+
+
+def get_model(params: AttributeDict) -> nn.Module:
+ encoder_embed = get_encoder_embed(params)
+ enc_asr = _build_zipformer("asr", params.output_downsampling_factor_asr, params)
+ enc_st = _build_zipformer("st", params.output_downsampling_factor_st, params)
+
+ decoder_asr = joiner_asr = None
+ decoder_st = joiner_st = None
+ if params.use_transducer:
+ decoder_asr = get_decoder(
+ "asr", params, params.vocab_size_asr, params.blank_id_asr
+ )
+ joiner_asr = get_joiner("asr", params, params.vocab_size_asr)
+ decoder_st = get_decoder("st", params, params.vocab_size_st, params.blank_id_st)
+ joiner_st = get_joiner("st", params, params.vocab_size_st)
+
+ attention_decoder_asr = attention_decoder_st = None
+ if params.use_attention_decoder:
+ attention_decoder_asr = get_attention_decoder(
+ "asr", params, params.vocab_size_asr, params.sos_id_asr, params.eos_id_asr
+ )
+ attention_decoder_st = get_attention_decoder(
+ "st", params, params.vocab_size_st, params.sos_id_st, params.eos_id_st
+ )
+
+ model = LCMASRTModel(
+ encoder_embed=encoder_embed,
+ enc_asr=enc_asr,
+ enc_st=enc_st,
+ decoder_asr=decoder_asr,
+ joiner_asr=joiner_asr,
+ attention_decoder_asr=attention_decoder_asr,
+ decoder_st=decoder_st,
+ joiner_st=joiner_st,
+ attention_decoder_st=attention_decoder_st,
+ encoder_dim_asr=max(_to_int_tuple(params.encoder_dim_asr)),
+ encoder_dim_st=max(_to_int_tuple(params.encoder_dim_st)),
+ decoder_dim_asr=params.decoder_dim_asr,
+ decoder_dim_st=params.decoder_dim_st,
+ vocab_size_asr=params.vocab_size_asr,
+ vocab_size_st=params.vocab_size_st,
+ output_downsampling_factor_asr=params.output_downsampling_factor_asr,
+ output_downsampling_factor_st=params.output_downsampling_factor_st,
+ num_srt_langs_asr=params.num_srt_langs_asr,
+ num_tgt_langs_ast=params.num_tgt_langs_ast,
+ num_experts_asr=params.num_experts_asr,
+ num_experts_ast=params.num_experts_ast,
+ entropy_reg_asr=params.entropy_reg_asr,
+ entropy_reg_ast=params.entropy_reg_ast,
+ temperature_asr=params.temperature_asr,
+ temperature_ast=params.temperature_ast,
+ asr_moe=params.asr_moe,
+ asr_src=params.asr_src,
+ ast_moe=params.ast_moe,
+ ast_tgt=params.ast_tgt,
+ use_transducer=params.use_transducer,
+ use_ctc_asr=params.use_ctc_asr,
+ use_ctc_st=params.use_ctc_st,
+ use_attention_decoder=params.use_attention_decoder,
+ freeze_asr=params.freeze_asr,
+ freeze_frontend=params.freeze_frontend,
+ )
+ return model
+
+
+def get_spec_augment(params: AttributeDict) -> SpecAugment:
+ num_frame_masks = int(10 * params.time_mask_ratio)
+ max_frames_mask_fraction = 0.15 * params.time_mask_ratio
+ logging.info(
+ f"num_frame_masks: {num_frame_masks}, max_frames_mask_fraction: {max_frames_mask_fraction}"
+ )
+ return SpecAugment(
+ time_warp_factor=0,
+ num_frame_masks=num_frame_masks,
+ features_mask_size=27,
+ num_feature_masks=2,
+ frames_mask_size=100,
+ max_frames_mask_fraction=max_frames_mask_fraction,
+ )
+
+
+# ------------------------------
+# Checkpoint helpers
+# ------------------------------
+
+
+def _load_model_weights_drop_st_head(
+ ckpt_path, model, st_prefixes=None, log_drops=True
+):
+ import logging
+
+ import torch
+
+ if st_prefixes is None:
+ st_prefixes = [
+ "decoder_st.",
+ "joiner_st.",
+ "simple_am_proj_st",
+ "simple_lm_proj_st",
+ "ctc_st",
+ "attention_decoder_st",
+ ]
+ ckpt = torch.load(ckpt_path, map_location="cpu")
+ src = ckpt["model"]
+ tgt_model = model.module if hasattr(model, "module") else model
+ tgt = tgt_model.state_dict()
+
+ def is_st_key(k: str) -> bool:
+ return any(k.startswith(p) or (p in k) for p in st_prefixes)
+
+ keep, dropped = {}, []
+ for k, v in src.items():
+ if k in tgt and not is_st_key(k):
+ keep[k] = v
+ else:
+ if is_st_key(k):
+ dropped.append(k)
+
+ missing, unexpected = tgt_model.load_state_dict(keep, strict=False)
+ if log_drops and dropped:
+ for k in dropped:
+ logging.warning(f"[CKPT] Drop ST param unconditionally: {k}")
+ rest = {
+ k: v
+ for k, v in ckpt.items()
+ if k not in ("model", "optimizer", "scheduler", "grad_scaler")
+ }
+ return {"missing": missing, "unexpected": unexpected, "dropped": dropped, **rest}
+
+
+def _safe_load_ckpt_with_filter(
+ ckpt_path: Path,
+ model: nn.Module,
+ drop_prefixes: Optional[List[str]] = None,
+ log_mismatch: bool = True,
+) -> Dict[str, Any]:
+ import logging
+
+ import torch
+
+ drop_prefixes = drop_prefixes or []
+ ckpt = torch.load(ckpt_path, map_location="cpu")
+ src = ckpt["model"]
+
+ tgt_model = model.module if hasattr(model, "module") else model
+ tgt_state = tgt_model.state_dict()
+
+ keep = {}
+ dropped = []
+
+ def _should_drop(name: str) -> bool:
+ return any(name.startswith(pfx) for pfx in drop_prefixes)
+
+ for k, v in src.items():
+ if _should_drop(k):
+ dropped.append(
+ (
+ k,
+ f"drop-by-prefix({[p for p in drop_prefixes if k.startswith(p)][0]})",
+ )
+ )
+ continue
+ if k in tgt_state and tgt_state[k].shape == v.shape:
+ keep[k] = v
+ else:
+ reason = "shape-mismatch" if (k in tgt_state) else "missing-in-target"
+ dropped.append((k, reason))
+
+ missing, unexpected = tgt_model.load_state_dict(keep, strict=False)
+
+ if log_mismatch:
+ for k, why in dropped:
+ old_shape = src[k].shape
+ new_shape = tgt_state[k].shape if k in tgt_state else "N/A"
+ logging.warning(f"[CKPT] Drop {k}: {why}; {old_shape} -> {new_shape}")
+ if missing:
+ logging.warning(
+ f"[CKPT] Missing keys after load (ok with strict=False): {missing[:8]}{' ...' if len(missing)>8 else ''}"
+ )
+ if unexpected:
+ logging.warning(
+ f"[CKPT] Unexpected keys in ckpt (ignored): {unexpected[:8]}{' ...' if len(unexpected)>8 else ''}"
+ )
+ rest = {k: v for k, v in ckpt.items() if k != "model"}
+ return rest
+
+
+def load_checkpoint_if_available(
+ params: AttributeDict,
+ model: nn.Module,
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+) -> Optional[Dict[str, Any]]:
+ if (
+ params.resume_from_checkpoint
+ and params.remove_st_head
+ and params.start_batch == 0
+ and params.start_epoch == 1
+ ):
+ filename = Path(params.resume_from_checkpoint)
+ assert filename.is_file(), f"{filename} does not exist!"
+ info = _load_model_weights_drop_st_head(filename, model)
+ if model_avg is not None:
+ with torch.no_grad():
+ model_avg.load_state_dict(
+ (model.module if hasattr(model, "module") else model).state_dict()
+ )
+ if "params" in info:
+ for k in [
+ "best_train_epoch",
+ "best_valid_epoch",
+ "batch_idx_train",
+ "best_train_loss",
+ "best_valid_loss",
+ "cur_epoch",
+ ]:
+ if k in info["params"]:
+ params[k] = info["params"][k]
+ return info
+
+ elif (
+ params.resume_from_checkpoint
+ and params.start_batch == 0
+ and params.start_epoch == 1
+ ):
+ filename = Path(params.resume_from_checkpoint)
+ elif params.start_batch > 0:
+ filename = params.exp_dir / f"checkpoint-{params.start_batch}.pt"
+ elif params.start_epoch > 1:
+ filename = params.exp_dir / f"epoch-{params.start_epoch-1}.pt"
+ else:
+ return None
+
+ assert filename.is_file(), f"{filename} does not exist!"
+
+ drop_prefixes = ["ast_moe_layer.router."]
+
+ rest = _safe_load_ckpt_with_filter(filename, model, drop_prefixes=drop_prefixes)
+
+ if model_avg is not None:
+ with torch.no_grad():
+ model_avg.load_state_dict(
+ (model.module if hasattr(model, "module") else model).state_dict()
+ )
+ saved_params: Dict[str, Any] = {}
+ reset_progress = getattr(params, "reset_progress_stats", False)
+ if params.resume_optimizer_scheduler_scaler and not reset_progress:
+ if "optimizer" in rest and optimizer is not None:
+ logging.info("Loading optimizer state dict")
+ try:
+ optimizer.load_state_dict(rest["optimizer"])
+ except Exception as e:
+ logging.warning(f"Load optimizer state failed: {e}")
+ if (
+ "scheduler" in rest
+ and scheduler is not None
+ and rest["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ try:
+ scheduler.load_state_dict(rest["scheduler"])
+ except Exception as e:
+ logging.warning(f"Load scheduler state failed: {e}")
+ if "grad_scaler" in rest:
+ saved_params["grad_scaler"] = rest["grad_scaler"]
+ elif params.resume_optimizer_scheduler_scaler and reset_progress:
+ logging.info(
+ "Skipping optimizer/scheduler/scaler restoration due to --reset-progress-stats=1"
+ )
+
+ if "sampler" in rest and not reset_progress:
+ saved_params["sampler"] = rest["sampler"]
+ elif "sampler" in rest and reset_progress:
+ logging.info(
+ "Skipping sampler state restoration due to --reset-progress-stats=1"
+ )
+
+ resume_state = rest.get("params")
+ if resume_state is None:
+ fallback_state = {k: rest[k] for k in _RESUME_STATE_KEYS if k in rest}
+ resume_state = fallback_state if fallback_state else None
+ if resume_state:
+ saved_params["params"] = resume_state
+ else:
+ saved_params["params"] = resume_state
+
+ if resume_state:
+ if not reset_progress:
+ for k in [
+ "best_train_epoch",
+ "best_valid_epoch",
+ "batch_idx_train",
+ "best_train_loss",
+ "best_valid_loss",
+ "cur_epoch",
+ ]:
+ if k in resume_state:
+ params[k] = resume_state[k]
+
+ if params.start_batch > 0 and "cur_epoch" in resume_state:
+ params["start_epoch"] = resume_state["cur_epoch"]
+
+ if "rng_state" in resume_state:
+ _restore_rng_state(resume_state["rng_state"])
+ else:
+ logging.info(
+ "reset-progress-stats=1: skipping training counters and RNG restoration."
+ )
+
+ return saved_params
+
+
+def save_checkpoint(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ model_avg: Optional[nn.Module] = None,
+ optimizer: Optional[torch.optim.Optimizer] = None,
+ scheduler: Optional[LRSchedulerType] = None,
+ sampler: Optional[CutSampler] = None,
+ scaler: Optional[GradScaler] = None,
+ rank: int = 0,
+) -> None:
+ if rank != 0:
+ return
+ params.rng_state = _capture_rng_state()
+ filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt"
+ params_payload = _build_params_payload(params)
+ save_checkpoint_impl(
+ filename=filename,
+ model=model,
+ model_avg=model_avg,
+ params=params_payload,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ if params.best_train_epoch == params.cur_epoch:
+ copyfile(src=filename, dst=params.exp_dir / "best-train-loss.pt")
+ if params.best_valid_epoch == params.cur_epoch:
+ copyfile(src=filename, dst=params.exp_dir / "best-valid-loss.pt")
+
+
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import torch
+
+
+def _normalize_lang_tag(tag: Optional[str]) -> Optional[str]:
+ if tag is None:
+ return None
+ if not isinstance(tag, str):
+ return tag
+ normalized = tag.strip()
+ if not normalized:
+ return None
+ return normalized.lower()
+
+
+def build_srctgt_lang_list(src_langs: List[str], tgt_langs: List[str]) -> List[str]:
+ """
+ Generate human-readable labels for every (src, tgt) language pair so
+ MoE routing stats can show directions like en->de or de->en.
+ """
+ if not src_langs or not tgt_langs:
+ return []
+ return [f"{src}->{tgt}" for src in src_langs for tgt in tgt_langs]
+
+
+def _extract_st_texts_and_lang_ids(
+ supervisions: List[Dict[str, Any]],
+ use_tgt: bool,
+ tgt_lang2id: Dict[str, int],
+ default_lang: str = None,
+):
+ default_lang = _normalize_lang_tag(default_lang)
+ st_texts = []
+ lang_ids: list[int] = []
+
+ for cut in supervisions["cut"]:
+ for supervision in cut.supervisions:
+ if hasattr(supervision, "custom") and "st_text" in supervision.custom:
+ lang_tag = None
+ if "lang" in supervision.custom and supervision.custom["lang"]:
+ lang_tag = _normalize_lang_tag(supervision.custom["lang"])
+ if lang_tag is None and default_lang is not None:
+ lang_tag = default_lang
+
+ if use_tgt:
+ if lang_tag is None:
+ raise ValueError(
+ "Missing custom['lang'] and no default_lang provided."
+ )
+ if lang_tag not in tgt_lang2id:
+ raise KeyError(
+ f"Unknown target language tag: {lang_tag}. "
+ f"Known: {list(tgt_lang2id.keys())}"
+ )
+
+ supervision.custom["st_text"] = (
+ f"<2{lang_tag}>" + supervision.custom["st_text"]
+ )
+ lang_ids.append(tgt_lang2id[lang_tag])
+ else:
+ if lang_tag is None:
+ lang_ids.append(0)
+ else:
+ lang_ids.append(tgt_lang2id.get(lang_tag, 0))
+
+ st_texts.append(supervision.custom["st_text"])
+ tgt_lang_ids = torch.tensor(lang_ids, dtype=torch.long)
+ return st_texts, tgt_lang_ids
+
+
+def asr_source_lang_tensor(
+ supervisions: Dict[str, Any],
+ srt_lang2id: Dict[str, int],
+ *,
+ strict: bool = True,
+) -> torch.LongTensor:
+ tags: List[Optional[str]] = []
+ cuts: Iterable[Any] = supervisions.get("cut", [])
+ for cut in cuts:
+ sups = getattr(cut, "supervisions", None)
+ if sups is None and isinstance(cut, dict):
+ sups = cut.get("supervisions", [])
+ if not sups:
+ continue
+
+ for sup in sups:
+ if isinstance(sup, dict):
+ text = sup.get("text")
+ lang = sup.get("language")
+ else:
+ text = getattr(sup, "text", None)
+ lang = getattr(sup, "language", None)
+
+ if text is None:
+ continue
+
+ lang = _normalize_lang_tag(lang)
+
+ if lang is None:
+ raise KeyError("Missing supervision['language'] for an ASR sample.")
+ tags.append(lang)
+
+ if "text" in supervisions and isinstance(supervisions["text"], list):
+ assert len(tags) == len(
+ supervisions["text"]
+ ), f"The number of ASR languages ({len(tags)}) is inconsistent with the number of texts ({len(supervisions['text'])})."
+ if strict:
+ ids = []
+ for t in tags:
+ if t not in srt_lang2id:
+ raise ValueError(
+ f"Unknown source language: {t}. Known: {list(srt_lang2id.keys())}"
+ )
+ ids.append(srt_lang2id[t])
+ else:
+ ids = [srt_lang2id.get(t, 0) for t in tags]
+
+ return torch.tensor(ids, dtype=torch.long)
+
+
+def _create_moe_stat_buffers(num_experts: int):
+ if num_experts <= 0:
+ return None, None
+ return (
+ defaultdict(lambda: torch.zeros(num_experts, dtype=torch.float64)),
+ defaultdict(int),
+ )
+
+
+def _accumulate_moe_stats(storage, counts, batch_info):
+ if storage is None or counts is None or batch_info is None:
+ return
+ lang_ids, weights = batch_info
+ if lang_ids is None or weights is None:
+ return
+ lang_list = lang_ids.tolist()
+ for idx, w in zip(lang_list, weights):
+ storage[idx] += w.to(storage[idx].dtype)
+ counts[idx] += 1
+
+
+def _log_moe_stats(
+ task_name: str,
+ storage,
+ counts,
+ lang_list: List[str],
+ global_step: int,
+ tb_writer: Optional[SummaryWriter] = None,
+):
+ if storage is None or counts is None or not storage:
+ return
+ logging.info("===== MoE routing stats (%s) =====", task_name)
+ for lang_id in sorted(storage.keys()):
+ count = counts[lang_id]
+ if count == 0:
+ continue
+ avg = storage[lang_id] / count
+ lang = lang_list[lang_id] if 0 <= lang_id < len(lang_list) else str(lang_id)
+ dist = ", ".join(f"e{i}:{float(val):.3f}" for i, val in enumerate(avg))
+ logging.info(" %s (id=%d, n=%d): %s", lang, lang_id, count, dist)
+ if tb_writer is not None:
+ for expert_idx, val in enumerate(avg):
+ tb_writer.add_scalar(
+ f"moe/{task_name}/{lang}/expert_{expert_idx}",
+ float(val),
+ global_step,
+ )
+
+
+def _reset_moe_stats(storage, counts):
+ if storage is not None:
+ storage.clear()
+ if counts is not None:
+ counts.clear()
+
+
+def compute_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ batch: dict,
+ is_training: bool,
+ spec_augment: Optional[SpecAugment] = None,
+) -> Tuple[Tensor, MetricsTracker]:
+ device = model.device if isinstance(model, DDP) else next(model.parameters()).device
+
+ feature = batch["inputs"].to(device) # (N, T, C)
+ assert feature.ndim == 3
+ use_tgt = params.use_tgt
+ supervisions = batch["supervisions"]
+ feature_lens = supervisions["num_frames"].to(device)
+
+ collect_moe_stats = bool(getattr(params, "dump_moe_routing_stats", False))
+
+ texts_asr: List[str] = supervisions["text"]
+ if params.asr_src:
+ srt_lang_ids = asr_source_lang_tensor(
+ supervisions, params.srt_lang2id, strict=True
+ )
+ else:
+ srt_lang_ids = None
+ stats_srt_lang_ids = srt_lang_ids
+
+ if params.enable_st:
+ texts_st, tgt_lang_ids = _extract_st_texts_and_lang_ids(
+ supervisions, use_tgt, params.tgt_lang2id
+ )
+ else:
+ texts_st, tgt_lang_ids = [], None
+
+ stats_tgt_lang_ids = tgt_lang_ids
+ if collect_moe_stats and params.use_cr_ctc:
+ if stats_srt_lang_ids is not None:
+ stats_srt_lang_ids = stats_srt_lang_ids.repeat(2)
+ if stats_tgt_lang_ids is not None:
+ stats_tgt_lang_ids = stats_tgt_lang_ids.repeat(2)
+
+ speace_id = sp_st.piece_to_id("▁")
+ y_asr = k2.RaggedTensor(sp_asr.encode(texts_asr, out_type=int))
+
+ if params.enable_st:
+ if use_tgt:
+ y_st_sp_list = sp_st.encode(texts_st, out_type=int)
+ y_st_sp = [
+ ids[1:] if ids and ids[0] == speace_id else ids for ids in y_st_sp_list
+ ]
+ y_st = k2.RaggedTensor(y_st_sp)
+ else:
+ y_st = k2.RaggedTensor(sp_st.encode(texts_st, out_type=int))
+ else:
+ y_st = None
+
+ # SpecAug/CR-CTC supports supervision segments
+ use_cr_ctc = params.use_cr_ctc
+ use_spec_aug = use_cr_ctc and is_training
+ if use_spec_aug:
+ sv = supervisions
+ supervision_segments = torch.stack(
+ [sv["sequence_idx"], sv["start_frame"], sv["num_frames"]], dim=1
+ )
+ else:
+ supervision_segments = None
+
+ batch_idx_train = params.batch_idx_train
+ warm_step = params.warm_step
+
+ with torch.set_grad_enabled(is_training):
+ outputs = model(
+ x=feature,
+ x_lens=feature_lens,
+ y_asr=y_asr,
+ y_st=y_st,
+ srt_lang_ids=srt_lang_ids,
+ tgt_lang_ids=tgt_lang_ids,
+ prune_range_asr=params.prune_range_asr,
+ prune_range_st=params.prune_range_st,
+ am_scale=params.am_scale,
+ lm_scale=params.lm_scale,
+ use_cr_ctc=use_cr_ctc,
+ use_spec_aug=use_spec_aug,
+ spec_augment=spec_augment,
+ supervision_segments=supervision_segments,
+ time_warp_factor=params.spec_aug_time_warp_factor,
+ enable_st=params.enable_st,
+ )
+
+ # Expected common return (allow extra tails):
+ # simple_asr, simple_st, pruned_asr, pruned_st, ctc_asr, ctc_st,
+ # attn_asr, attn_st, cr_asr, cr_st, [optional distill]
+ (
+ simple_asr,
+ simple_st,
+ pruned_asr,
+ pruned_st,
+ ctc_asr,
+ ctc_st,
+ attn_asr,
+ attn_st,
+ cr_asr,
+ cr_st,
+ moe_ent_loss,
+ *rest,
+ ) = outputs
+
+ # distill = rest[0] if (len(rest) > 0 and isinstance(rest[0], torch.Tensor)) else None
+
+ loss = torch.tensor(0.0, device=feature.device)
+ loss_asr = torch.tensor(0.0, device=feature.device)
+ loss_st = torch.tensor(0.0, device=feature.device)
+
+ # Transducer head (simple + pruned)
+ if params.use_transducer:
+
+ def _warm_scale(s: float) -> Tuple[float, float]:
+ simple_scale = (
+ s
+ if batch_idx_train >= warm_step
+ else 1.0 - (batch_idx_train / warm_step) * (1.0 - s)
+ )
+ pruned_scale = (
+ 1.0
+ if batch_idx_train >= warm_step
+ else 0.1 + 0.9 * (batch_idx_train / warm_step)
+ )
+ return simple_scale, pruned_scale
+
+ s_asr_scale, p_asr_scale = _warm_scale(params.simple_loss_scale)
+ loss_asr += s_asr_scale * simple_asr + p_asr_scale * pruned_asr
+
+ if params.enable_st:
+ s_st_scale, p_st_scale = _warm_scale(params.simple_loss_scale)
+ loss_st += s_st_scale * simple_st + p_st_scale * pruned_st
+ # loss = loss + params.task_weight_asr * (s_asr * simple_asr + p_asr * pruned_asr)
+ # loss = loss + params.task_weight_st * (s_st * simple_st + p_st * pruned_st)
+
+ # CTC (with optional CR-CTC)
+ if params.use_ctc_asr:
+ # loss = loss + params.task_weight_asr * (params.ctc_loss_scale * ctc_asr)
+ # loss = loss + params.task_weight_st * (params.ctc_loss_scale * ctc_st)
+ loss_asr += params.ctc_loss_scale * ctc_asr # ctc_loss_scale default 0.2
+ if params.use_cr_ctc:
+ # loss = loss + params.task_weight_asr * (params.cr_loss_scale * cr_asr)
+ # loss = loss + params.task_weight_st * (params.cr_loss_scale * cr_st)
+ loss_asr += params.cr_loss_scale * cr_asr # cr_loss_scale default 0.2
+
+ if params.use_ctc_st and params.enable_st:
+ # loss = loss + params.task_weight_asr * (params.ctc_loss_scale * ctc_asr)
+ # loss = loss + params.task_weight_st * (params.ctc_loss_scale * ctc_st)
+ loss_st += params.ctc_loss_scale * ctc_st
+ if params.use_cr_ctc:
+ # loss = loss + params.task_weight_asr * (params.cr_loss_scale * cr_asr)
+ # loss = loss + params.task_weight_st * (params.cr_loss_scale * cr_st)
+ loss_st += params.cr_loss_scale * cr_st
+
+ # Attention decoder
+ if params.use_attention_decoder:
+ # loss = loss + params.task_weight_asr * (params.attention_decoder_loss_scale * attn_asr)
+ # loss = loss + params.task_weight_st * (params.attention_decoder_loss_scale * attn_st)
+ loss_asr += params.attention_decoder_loss_scale * attn_asr
+ if params.enable_st:
+ loss_st += params.attention_decoder_loss_scale * attn_st
+
+ # Optional self-distillation term
+ # if distill is not None and params.self_distill_scale > 0:
+ # loss = loss + params.self_distill_scale * distill
+
+ # loss = params.task_weight_asr * loss_asr + params.task_weight_st * loss_st + moe_ent_loss
+ if params.enable_st:
+ loss = (
+ params.task_weight_asr * loss_asr
+ + params.task_weight_st * loss_st
+ + moe_ent_loss
+ )
+ else:
+ loss = params.task_weight_asr * loss_asr + moe_ent_loss
+
+ assert loss.requires_grad == is_training
+
+ # Metrics
+ info = MetricsTracker()
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ info["frames"] = (feature_lens // params.subsampling_factor).sum().item()
+
+ info["loss"] = loss.detach().cpu().item()
+ info["loss_asr"] = loss_asr.detach().cpu().item()
+ info["loss_ast"] = loss_st.detach().cpu().item()
+ info["moe_ent_loss"] = moe_ent_loss.detach().cpu().item()
+ if params.use_transducer:
+ info["simple_asr"] = simple_asr.detach().cpu().item()
+ info["pruned_asr"] = pruned_asr.detach().cpu().item()
+
+ if params.use_transducer and params.enable_st:
+ info["simple_ast"] = simple_st.detach().cpu().item()
+ info["pruned_ast"] = pruned_st.detach().cpu().item()
+
+ if params.use_ctc_asr:
+ info["ctc_asr"] = ctc_asr.detach().cpu().item()
+ if params.use_cr_ctc:
+ info["cr_asr"] = cr_asr.detach().cpu().item()
+ # if params.use_ctc_st:
+ # info["ctc_st"] = ctc_st.detach().cpu().item()
+ # if params.use_cr_ctc:
+ # info["cr_st"] = cr_st.detach().cpu().item()
+ if params.use_ctc_st and params.enable_st:
+ info["ctc_ast"] = ctc_st.detach().cpu().item()
+ if params.use_cr_ctc:
+ info["cr_ast"] = cr_st.detach().cpu().item()
+
+ if params.use_attention_decoder:
+ info["attn_asr"] = attn_asr.detach().cpu().item()
+
+ if params.use_attention_decoder and params.enable_st:
+ info["attn_ast"] = attn_st.detach().cpu().item()
+ # info["attn_st"] = attn_st.detach().cpu().item()
+ # if distill is not None:
+ # info["distill"] = distill.detach().cpu().item()
+
+ moe_batch_stats = None
+ if collect_moe_stats:
+ moe_batch_stats = {}
+ base_model = model.module if isinstance(model, DDP) else model
+ if hasattr(base_model, "asr_moe_layer"):
+ weights_asr = getattr(base_model.asr_moe_layer, "last_router_weights", None)
+ if (
+ weights_asr is not None
+ and stats_srt_lang_ids is not None
+ and weights_asr.ndim == 3
+ ):
+ mean_asr = weights_asr.detach().mean(dim=0).to(torch.float64).cpu()
+ if mean_asr.size(0) == stats_srt_lang_ids.numel():
+ moe_batch_stats["asr"] = (
+ stats_srt_lang_ids.detach().cpu(),
+ mean_asr,
+ )
+ else:
+ logging.warning(
+ f"({mean_asr.size(0)} vs {stats_srt_lang_ids.numel()})"
+ )
+ if params.enable_st and hasattr(base_model, "ast_moe_layer"):
+ weights_st = getattr(base_model.ast_moe_layer, "last_router_weights", None)
+ if (
+ weights_st is not None
+ and stats_tgt_lang_ids is not None
+ and weights_st.ndim == 3
+ ):
+ mean_st = weights_st.detach().mean(dim=0).to(torch.float64).cpu()
+ if mean_st.size(0) == stats_tgt_lang_ids.numel():
+ moe_batch_stats["st"] = (
+ stats_tgt_lang_ids.detach().cpu(),
+ mean_st,
+ )
+ else:
+ logging.warning(
+ f"({mean_st.size(0)} vs {stats_tgt_lang_ids.numel()})"
+ )
+ if not moe_batch_stats:
+ moe_batch_stats = None
+
+ return loss, info, moe_batch_stats
+
+
+def compute_validation_loss(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ valid_dl: torch.utils.data.DataLoader,
+ world_size: int = 1,
+) -> MetricsTracker:
+ if dist.is_initialized() and dist.get_rank() != 0:
+ return MetricsTracker()
+
+ model.eval()
+ tot_loss = MetricsTracker()
+ collect_moe_stats = bool(getattr(params, "dump_moe_routing_stats", False))
+ st_lang_labels = params.tgt_lang_list
+ moe_stats_asr = moe_counts_asr = moe_stats_st = moe_counts_st = None
+ if collect_moe_stats:
+ base_model = model.module if isinstance(model, DDP) else model
+ asr_experts = getattr(
+ getattr(base_model, "asr_moe_layer", None), "num_experts", 0
+ )
+ ast_experts = getattr(
+ getattr(base_model, "ast_moe_layer", None), "num_experts", 0
+ )
+ moe_stats_asr, moe_counts_asr = _create_moe_stat_buffers(asr_experts)
+ moe_stats_st, moe_counts_st = _create_moe_stat_buffers(ast_experts)
+ for batch_idx, batch in enumerate(valid_dl):
+ loss, info, batch_moe_stats = compute_loss(
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ batch=batch,
+ is_training=False,
+ )
+ if collect_moe_stats and batch_moe_stats:
+ _accumulate_moe_stats(
+ moe_stats_asr,
+ moe_counts_asr,
+ batch_moe_stats.get("asr"),
+ )
+ _accumulate_moe_stats(
+ moe_stats_st,
+ moe_counts_st,
+ batch_moe_stats.get("st"),
+ )
+ assert loss.requires_grad is False
+ tot_loss = tot_loss + info
+
+ loss_value = tot_loss["loss"] / max(tot_loss["frames"], 1)
+ if loss_value < params.best_valid_loss:
+ params.best_valid_epoch = params.cur_epoch
+ params.best_valid_loss = loss_value
+ if collect_moe_stats:
+ _log_moe_stats(
+ "valid_ASR",
+ moe_stats_asr,
+ moe_counts_asr,
+ params.srt_lang_list,
+ params.batch_idx_train,
+ None,
+ )
+ _log_moe_stats(
+ "valid_ST",
+ moe_stats_st,
+ moe_counts_st,
+ st_lang_labels,
+ params.batch_idx_train,
+ None,
+ )
+ _reset_moe_stats(moe_stats_asr, moe_counts_asr)
+ _reset_moe_stats(moe_stats_st, moe_counts_st)
+ return tot_loss
+
+
+def _reapply_freeze_asr(model):
+ m = model.module if isinstance(model, DDP) else model
+ if hasattr(m, "freeze_asr") and m.freeze_asr and hasattr(m, "_apply_freeze_asr"):
+ m._apply_freeze_asr()
+
+
+# ------------------------------
+# Training loop
+# ------------------------------
+
+
+def train_one_epoch(
+ params: AttributeDict,
+ model: Union[nn.Module, DDP],
+ optimizer: torch.optim.Optimizer,
+ scheduler: LRSchedulerType,
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ train_dl: torch.utils.data.DataLoader,
+ valid_dl: torch.utils.data.DataLoader,
+ scaler: GradScaler,
+ spec_augment: Optional[SpecAugment] = None,
+ model_avg: Optional[nn.Module] = None,
+ tb_writer: Optional[SummaryWriter] = None,
+ world_size: int = 1,
+ rank: int = 0,
+) -> None:
+ model.train()
+ _reapply_freeze_asr(model)
+ tot_loss = MetricsTracker()
+ saved_bad_model = False
+ collect_moe_stats = (
+ bool(getattr(params, "dump_moe_routing_stats", False)) and rank == 0
+ )
+ st_lang_labels = params.tgt_lang_list
+ moe_stats_asr = moe_counts_asr = moe_stats_st = moe_counts_st = None
+ if collect_moe_stats:
+ base_model = model.module if isinstance(model, DDP) else model
+ asr_experts = getattr(
+ getattr(base_model, "asr_moe_layer", None), "num_experts", 0
+ )
+ ast_experts = getattr(
+ getattr(base_model, "ast_moe_layer", None), "num_experts", 0
+ )
+ moe_stats_asr, moe_counts_asr = _create_moe_stat_buffers(asr_experts)
+ moe_stats_st, moe_counts_st = _create_moe_stat_buffers(ast_experts)
+
+ def save_bad_model(suffix: str = ""):
+ params.rng_state = _capture_rng_state()
+ params_payload = _build_params_payload(params)
+ save_checkpoint_impl(
+ filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt",
+ model=model.module if isinstance(model, DDP) else model,
+ model_avg=model_avg,
+ params=params_payload,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=0,
+ )
+
+ for batch_idx, batch in enumerate(train_dl):
+ if batch_idx % 10 == 0:
+ set_batch_count(model, get_adjusted_batch_count(params))
+ params.batch_idx_train += 1
+ batch_size = len(batch["supervisions"]["text"])
+
+ try:
+ with torch.cuda.amp.autocast(
+ enabled=params.use_autocast, dtype=params.dtype
+ ):
+ loss, loss_info, batch_moe_stats = compute_loss(
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ batch=batch,
+ is_training=True,
+ spec_augment=spec_augment,
+ )
+ tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info
+ if collect_moe_stats and batch_moe_stats:
+ _accumulate_moe_stats(
+ moe_stats_asr,
+ moe_counts_asr,
+ batch_moe_stats.get("asr"),
+ )
+ _accumulate_moe_stats(
+ moe_stats_st,
+ moe_counts_st,
+ batch_moe_stats.get("st"),
+ )
+
+ scaler.scale(loss).backward()
+ scheduler.step_batch(params.batch_idx_train)
+ scaler.step(optimizer)
+ scaler.update()
+ optimizer.zero_grad()
+ except Exception as e:
+ logging.info(f"Caught exception: {e}.")
+ save_bad_model()
+ display_and_save_batch(batch, params=params, sp_asr=sp_asr, sp_st=sp_st)
+ raise
+
+ if params.print_diagnostics and batch_idx == 5:
+ return
+
+ if (
+ rank == 0
+ and params.batch_idx_train > 0
+ and params.batch_idx_train % params.average_period == 0
+ ):
+ update_averaged_model(
+ params=params,
+ model_cur=model.module if isinstance(model, DDP) else model,
+ model_avg=model_avg,
+ )
+
+ if (
+ params.batch_idx_train > 0
+ and params.batch_idx_train % params.save_every_n == 0
+ ):
+ params.rng_state = _capture_rng_state()
+ params_payload = _build_params_payload(params)
+ save_checkpoint_with_global_batch_idx(
+ out_dir=params.exp_dir,
+ global_batch_idx=params.batch_idx_train,
+ model=model.module if isinstance(model, DDP) else model,
+ model_avg=model_avg,
+ params=params_payload,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+ remove_checkpoints(
+ out_dir=params.exp_dir, topk=params.keep_last_k, rank=rank
+ )
+
+ if params.use_autocast:
+ cur_grad_scale = scaler._scale.item()
+ if cur_grad_scale < 0.01:
+ if not saved_bad_model:
+ save_bad_model(suffix="-first-warning")
+ saved_bad_model = True
+ if not params.inf_check:
+ register_inf_check_hooks(model)
+ logging.warning(f"Grad scale is small: {cur_grad_scale}")
+ if cur_grad_scale < 1.0e-05:
+ save_bad_model()
+ raise_grad_scale_is_too_small_error(cur_grad_scale)
+ if (
+ batch_idx % 25 == 0
+ and cur_grad_scale < 2.0
+ or batch_idx % 100 == 0
+ and cur_grad_scale < 8.0
+ or batch_idx % 400 == 0
+ and cur_grad_scale < 32.0
+ ):
+ scaler.update(cur_grad_scale * 2.0)
+
+ if batch_idx % params.log_interval == 0:
+ cur_lr = max(scheduler.get_last_lr())
+ cur_grad_scale = scaler._scale.item() if params.use_autocast else 1.0
+ logging.info(
+ f"Epoch {params.cur_epoch}, batch {batch_idx}, loss[{loss_info}], tot_loss[{tot_loss}], "
+ f"batch size: {batch_size}, lr: {cur_lr:.2e}, "
+ + (f"grad_scale: {cur_grad_scale}" if params.use_autocast else "")
+ )
+ if tb_writer is not None:
+ tb_writer.add_scalar(
+ "train/learning_rate", cur_lr, params.batch_idx_train
+ )
+ loss_info.write_summary(
+ tb_writer, "train/current_", params.batch_idx_train
+ )
+ tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train)
+ if params.use_autocast:
+ tb_writer.add_scalar(
+ "train/grad_scale", cur_grad_scale, params.batch_idx_train
+ )
+ if collect_moe_stats:
+ _log_moe_stats(
+ "train_ASR",
+ moe_stats_asr,
+ moe_counts_asr,
+ params.srt_lang_list,
+ params.batch_idx_train,
+ tb_writer,
+ )
+ _log_moe_stats(
+ "train_ST",
+ moe_stats_st,
+ moe_counts_st,
+ st_lang_labels,
+ params.batch_idx_train,
+ tb_writer,
+ )
+ _reset_moe_stats(moe_stats_asr, moe_counts_asr)
+ _reset_moe_stats(moe_stats_st, moe_counts_st)
+
+ if batch_idx % params.valid_interval == 0 and not params.print_diagnostics:
+ logging.info("Computing validation loss")
+ valid_info = compute_validation_loss(
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ valid_dl=valid_dl,
+ world_size=world_size,
+ )
+ model.train()
+ _reapply_freeze_asr(model) # Freeze ASR
+ logging.info(f"Epoch {params.cur_epoch}, validation: {valid_info}")
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+ if tb_writer is not None:
+ valid_info.write_summary(
+ tb_writer, "train/valid_", params.batch_idx_train
+ )
+
+ loss_value = tot_loss["loss"] / max(tot_loss["frames"], 1)
+ params.train_loss = loss_value
+ if params.train_loss < params.best_train_loss:
+ params.best_train_epoch = params.cur_epoch
+ params.best_train_loss = params.train_loss
+
+ if collect_moe_stats:
+ _log_moe_stats(
+ "train_ASR",
+ moe_stats_asr,
+ moe_counts_asr,
+ params.srt_lang_list,
+ params.batch_idx_train,
+ tb_writer,
+ )
+ _log_moe_stats(
+ "train_ST",
+ moe_stats_st,
+ moe_counts_st,
+ st_lang_labels,
+ params.batch_idx_train,
+ tb_writer,
+ )
+ _reset_moe_stats(moe_stats_asr, moe_counts_asr)
+ _reset_moe_stats(moe_stats_st, moe_counts_st)
+
+
+# ------------------------------
+# Misc helpers
+# ------------------------------
+
+
+def display_and_save_batch(
+ batch: dict,
+ params: AttributeDict,
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+) -> None:
+ from lhotse.utils import uuid4
+
+ filename = f"{params.exp_dir}/batch-{uuid4()}.pt"
+ logging.info(f"Saving batch to {filename}")
+ torch.save(batch, filename)
+ use_tgt = params.use_tgt
+ supervisions = batch["supervisions"]
+ features = batch["inputs"]
+ logging.info(f"features shape: {features.shape}")
+
+ texts: List[str] = supervisions["text"]
+ if params.asr_src:
+ srt_lang_ids = asr_source_lang_tensor(
+ supervisions, params.srt_lang2id, strict=True
+ )
+ else:
+ srt_lang_ids = None
+
+ if params.enable_st:
+ texts_st, tgt_lang_ids = _extract_st_texts_and_lang_ids(
+ supervisions, use_tgt, params.tgt_lang2id
+ )
+ else:
+ texts_st, tgt_lang_ids = [], None
+
+ print("---" * 15)
+ print("DEBUG: Displaying text from the failing batch:")
+ if texts:
+ print(f" ASR Text[0]: {texts[0]}")
+ if srt_lang_ids is not None and len(srt_lang_ids):
+ print(f" ASR LangID[0]: {int(srt_lang_ids[0])}")
+ if texts_st:
+ print(f" ST Text[0]: {texts_st[0]}")
+ if len(tgt_lang_ids):
+ print(f" ST LangID[0]: {int(tgt_lang_ids[0])}")
+ print("---" * 15, flush=True)
+
+ y_asr = sp_asr.encode(texts, out_type=int)
+ y_st = sp_st.encode(texts_st, out_type=int)
+ num_tokens = sum(len(i) for i in y_asr) + sum(len(i) for i in y_st)
+ logging.info(f"num tokens (ASR+ST): {num_tokens}")
+
+
+def scan_pessimistic_batches_for_oom(
+ model: Union[nn.Module, DDP],
+ train_dl: torch.utils.data.DataLoader,
+ optimizer: torch.optim.Optimizer,
+ sp_asr: spm.SentencePieceProcessor,
+ sp_st: spm.SentencePieceProcessor,
+ params: AttributeDict,
+ spec_augment: Optional[SpecAugment] = None,
+):
+ from lhotse.dataset import find_pessimistic_batches
+
+ logging.info(
+ "Sanity check -- see if any of the batches in epoch 1 would cause OOM."
+ )
+ batches, crit_values = find_pessimistic_batches(train_dl.sampler)
+ for criterion, cuts in batches.items():
+ batch = train_dl.dataset[cuts]
+ try:
+ with torch.cuda.amp.autocast(
+ enabled=params.use_autocast, dtype=params.dtype
+ ):
+ loss, _, _ = compute_loss(
+ params=params,
+ model=model,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ batch=batch,
+ is_training=True,
+ spec_augment=spec_augment,
+ )
+ loss.backward()
+ optimizer.zero_grad()
+ except Exception as e:
+ if "CUDA out of memory" in str(e):
+ logging.error(
+ "Your GPU ran out of memory with the current max_duration setting. "
+ "Decrease max_duration and try again.\n"
+ f"Failing criterion: {criterion} (= {crit_values[criterion]}) ..."
+ )
+ display_and_save_batch(batch, params=params, sp_asr=sp_asr, sp_st=sp_st)
+ raise
+ logging.info(
+ f"Maximum memory allocated so far is {torch.cuda.max_memory_allocated()//1000000}MB"
+ )
+
+
+# ------------------------------
+# Entry point (multi-GPU safe)
+# ------------------------------
+
+
+def get_rank_info() -> Tuple[int, int]:
+ if "RANK" in os.environ and "LOCAL_RANK" in os.environ:
+ rank = int(os.environ["RANK"])
+ local_rank = int(os.environ["LOCAL_RANK"])
+ elif "SLURM_PROCID" in os.environ and "SLURM_LOCALID" in os.environ:
+ rank = int(os.environ["SLURM_PROCID"])
+ local_rank = int(os.environ["SLURM_LOCALID"])
+ else:
+ rank = 0
+ local_rank = 0
+ return rank, local_rank
+
+
+def setup_cuda_device(local_rank: int) -> None:
+ device_count = torch.cuda.device_count()
+ if device_count == 1:
+ torch.cuda.set_device(0)
+ elif local_rank < device_count:
+ torch.cuda.set_device(local_rank)
+ else:
+ raise RuntimeError(
+ f"[setup_cuda_device] local_rank={local_rank} exceeds visible CUDA devices ({device_count})"
+ )
+ print(f"[setup_cuda_device] Using CUDA:{torch.cuda.current_device()}")
+
+
+def run(rank: int, world_size: int, args: argparse.Namespace) -> None:
+ params = get_params()
+ params.update(vars(args))
+ if rank != 0:
+ params.dump_moe_routing_stats = False
+
+ # Alias encoder param names for convenience (ASR/ST)
+ # Convert CLI into attribute names used by builders
+ for name in (
+ "num_encoder_layers",
+ "downsampling_factor",
+ "feedforward_dim",
+ "num_heads",
+ "encoder_dim",
+ "query_head_dim",
+ "value_head_dim",
+ "pos_head_dim",
+ "pos_dim",
+ "encoder_unmasked_dim",
+ "cnn_module_kernel",
+ ):
+ setattr(params, f"{name}_asr", getattr(params, f"{name}_asr"))
+ setattr(params, f"{name}_st", getattr(params, f"{name}_st"))
+
+ fix_random_seed(params.seed)
+
+ local_rank = int(os.environ.get("LOCAL_RANK", 0))
+ if world_size > 1:
+ setup_dist(
+ rank=rank,
+ world_size=world_size,
+ master_port=params.master_port,
+ local_rank=local_rank,
+ )
+
+ setup_logger(f"{params.exp_dir}/log/log-train")
+ logging.info("Training started")
+
+ tb_writer = (
+ SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard")
+ if args.tensorboard and rank == 0
+ else None
+ )
+
+ device = (
+ torch.device("cuda", local_rank)
+ if torch.cuda.is_available()
+ else torch.device("cpu")
+ )
+ logging.info(f"Device: {device}")
+
+ # Tokenizers
+ sp_asr = spm.SentencePieceProcessor()
+ sp_asr.load(params.bpe_model_asr)
+ sp_st = spm.SentencePieceProcessor()
+ sp_st.load(params.bpe_model_st)
+
+ # Ids and vocab sizes per task
+ params.blank_id_asr = sp_asr.piece_to_id("")
+ params.sos_id_asr = params.eos_id_asr = sp_asr.piece_to_id("")
+ params.vocab_size_asr = sp_asr.get_piece_size()
+
+ params.blank_id_st = sp_st.piece_to_id("")
+ params.sos_id_st = params.eos_id_st = sp_st.piece_to_id("")
+ params.vocab_size_st = sp_st.get_piece_size()
+
+ # sp_st = SentencePieceProcessor() already loaded
+ params.tgt_lang_list = [s.strip() for s in params.tgt_langs.split(",") if s.strip()]
+ params.tgt_lang2id = {lg: i for i, lg in enumerate(params.tgt_lang_list)}
+ params.num_tgt_langs_ast = len(params.tgt_lang_list)
+
+ params.srt_lang_list = [s.strip() for s in params.srt_langs.split(",") if s.strip()]
+ params.srt_lang2id = {lg: i for i, lg in enumerate(params.srt_lang_list)}
+ params.num_srt_langs_asr = len(params.srt_lang_list)
+ params.srctgt_lang_list = build_srctgt_lang_list(
+ params.srt_lang_list, params.tgt_lang_list
+ )
+ # print(f"num_srt_langs_asr:{params.num_srt_langs_asr}")
+
+ # AMP dtype
+ if params.use_bf16:
+ assert torch.cuda.is_bf16_supported(), "Your GPU does not support bf16!"
+ assert not params.use_fp16, "Use either fp16 or bf16, not both."
+ params.dtype = torch.bfloat16
+ params.use_autocast = True
+ elif params.use_fp16:
+ params.dtype = torch.float16
+ params.use_autocast = True
+ else:
+ params.dtype = torch.float32
+ params.use_autocast = False
+ logging.info(f"Using dtype={params.dtype}; AMP={params.use_autocast}")
+
+ logging.info(params)
+ logging.info("About to create model")
+ model = get_model(params)
+ num_param = sum(p.numel() for p in model.parameters())
+ logging.info(f"Number of model parameters: {num_param}")
+
+ if params.use_cr_ctc:
+ assert params.use_ctc_asr
+ assert not params.enable_spec_aug # we will do spec_augment in model.py
+ spec_augment = get_spec_augment(params)
+ else:
+ spec_augment = None
+
+ assert params.save_every_n >= params.average_period
+ model_avg: Optional[nn.Module] = None
+ if rank == 0:
+ model_avg = copy.deepcopy(model).to(torch.float64)
+
+ assert params.start_epoch > 0
+
+ checkpoints = load_checkpoint_if_available(
+ params=params, model=model, model_avg=model_avg
+ )
+ _reapply_freeze_asr(model) # Freeze ASR
+ # logging.info(f"{model}")
+ model.to(device)
+ if world_size > 1:
+ logging.info("Using DDP")
+ model = DDP(model, device_ids=[local_rank], find_unused_parameters=True)
+
+ optimizer = ScaledAdam(
+ get_parameter_groups_with_lrs(model, lr=params.base_lr, include_names=True),
+ lr=params.base_lr,
+ clipping_scale=2.0,
+ )
+ scheduler: LRSchedulerType = Eden(
+ optimizer, params.lr_batches, params.lr_epochs, warmup_start=0.1
+ )
+
+ if params.resume_optimizer_scheduler_scaler:
+ if checkpoints and "optimizer" in checkpoints:
+ logging.info("Loading optimizer state dict")
+ optimizer.load_state_dict(checkpoints["optimizer"])
+ if (
+ checkpoints
+ and "scheduler" in checkpoints
+ and checkpoints["scheduler"] is not None
+ ):
+ logging.info("Loading scheduler state dict")
+ scheduler.load_state_dict(checkpoints["scheduler"])
+
+ if params.print_diagnostics:
+ opts = diagnostics.TensorDiagnosticOptions(512)
+ diagnostic = diagnostics.attach_diagnostics(model, opts)
+ if params.inf_check:
+ register_inf_check_hooks(model)
+
+ # Data
+ librispeech = LibriSpeechAsrDataModule(args)
+
+ if params.full_libri:
+ if args.train_cuts_paths:
+ train_cuts = load_cuts_lazy(args.train_cuts_paths, shuffle=True)
+ else:
+ train_cuts = librispeech.train_all_shuf_cuts()
+ else:
+ train_cuts = librispeech.train_clean_100_cuts()
+
+ def remove_short_and_long_utt(c: Cut) -> bool:
+ if (
+ c.duration < args.utterance_min_duration
+ or c.duration > args.utterance_max_duration
+ ):
+ return False
+ if getattr(c, "num_frames", None) is None:
+ print(
+ f"WARNING: Exclude cut {c.id}. num_frames is None (features not attached).",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+ supervision = c.supervisions[0]
+ custom = getattr(supervision, "custom", {}) or {}
+
+ text = getattr(supervision, "text", None)
+ if not text or not text.strip():
+ print(
+ f"WARNING: Exclude cut {c.id}. Missing or empty text field.",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+
+ st_text = custom.get("st_text")
+ if not st_text or not st_text.strip():
+ print(
+ f"WARNING: Exclude cut {c.id}. Missing or empty st_text field.",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+
+ language = getattr(supervision, "language", None)
+ if not language or not language.strip():
+ print(
+ f"WARNING: Exclude cut {c.id}. Missing or empty language field.",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+
+ lang = custom.get("lang")
+ if not lang or not lang.strip():
+ print(
+ f"WARNING: Exclude cut {c.id}. Missing or empty lang field.",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+ T = ((c.num_frames - 7) // 2 + 1) // 2
+
+ asr_tokens = sp_asr.encode(text, out_type=str)
+ asr_len = len(asr_tokens)
+
+ st_tokens = sp_st.encode(st_text, out_type=str)
+ st_len = len(st_tokens)
+
+ need = max(asr_len, st_len)
+ if T < need:
+ print(
+ f"WARNING: Exclude cut {c.id}. Frames(after)={T}. "
+ f"ASR tokens={asr_len}, ST tokens={st_len}",
+ file=sys.stderr,
+ flush=True,
+ )
+ return False
+ return True
+
+ train_cuts = train_cuts.filter(remove_short_and_long_utt)
+
+ if params.start_batch > 0 and checkpoints and "sampler" in checkpoints:
+ sampler_state_dict = checkpoints["sampler"]
+ else:
+ sampler_state_dict = None
+
+ train_dl = librispeech.train_dataloaders(
+ train_cuts, sampler_state_dict=sampler_state_dict
+ )
+ resume_mid_epoch = params.start_batch > 0 and sampler_state_dict is not None
+ skipped_initial_sampler_epoch = False
+
+ if args.valid_cuts_paths:
+ valid_cuts = load_cuts_lazy(args.valid_cuts_paths, shuffle=False)
+ else:
+ valid_cuts = librispeech.dev_clean_cuts()
+ valid_cuts += librispeech.dev_other_cuts()
+
+ valid_dl = librispeech.valid_dataloaders_gpt41(valid_cuts)
+
+ # if not params.skip_sanity_check and not params.print_diagnostics:
+ # scan_pessimistic_batches_for_oom(
+ # model=model,
+ # train_dl=train_dl,
+ # optimizer=optimizer,
+ # sp_asr=sp_asr,
+ # sp_st=sp_st,
+ # params=params,
+ # spec_augment=spec_augment,
+ # )
+
+ scaler = GradScaler(enabled=params.use_autocast, init_scale=1.0)
+ if params.resume_optimizer_scheduler_scaler:
+ if checkpoints and "grad_scaler" in checkpoints:
+ logging.info("Loading grad scaler state dict")
+ scaler.load_state_dict(checkpoints["grad_scaler"])
+
+ for epoch in range(params.start_epoch, params.num_epochs + 1):
+ scheduler.step_epoch(epoch - 1)
+ fix_random_seed(params.seed + epoch - 1)
+ skip_sampler_epoch = (
+ resume_mid_epoch
+ and not skipped_initial_sampler_epoch
+ and epoch == params.start_epoch
+ )
+ if skip_sampler_epoch:
+ logging.info(
+ "Skipping sampler.set_epoch for resumed epoch; sampler state restored from checkpoint."
+ )
+ skipped_initial_sampler_epoch = True
+ else:
+ train_dl.sampler.set_epoch(epoch - 1)
+
+ if tb_writer is not None:
+ tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train)
+ params.cur_epoch = epoch
+
+ train_one_epoch(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sp_asr=sp_asr,
+ sp_st=sp_st,
+ train_dl=train_dl,
+ valid_dl=valid_dl,
+ scaler=scaler,
+ spec_augment=spec_augment,
+ tb_writer=tb_writer,
+ world_size=world_size,
+ rank=rank,
+ )
+
+ if params.print_diagnostics:
+ diagnostic.print_diagnostics()
+ break
+
+ save_checkpoint(
+ params=params,
+ model=model,
+ model_avg=model_avg,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ sampler=train_dl.sampler,
+ scaler=scaler,
+ rank=rank,
+ )
+
+ logging.info("Done!")
+
+ if world_size > 1:
+ torch.distributed.barrier()
+ cleanup_dist()
+
+
+def main() -> None:
+ parser = get_parser()
+ args = parser.parse_args()
+ args.exp_dir = Path(args.exp_dir)
+
+ logging.info("## get rank information ...")
+ rank, local_rank = get_rank_info()
+ world_size = args.world_size
+
+ setup_cuda_device(local_rank)
+ run(rank=rank, world_size=world_size, args=args)
+
+ if dist.is_available() and dist.is_initialized():
+ dist.destroy_process_group()
+
+
+torch.set_num_threads(1)
+torch.set_num_interop_threads(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc.sh b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc.sh
new file mode 100644
index 0000000000..d713ff2f82
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc.sh
@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=900
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=0.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=0
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --enable-st 0 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_moe.sh b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_moe.sh
new file mode 100644
index 0000000000..d37db10403
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_moe.sh
@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=900
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=0.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=0
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --enable-st 0 \
+ --asr-moe 1 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 8 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_s_bias.sh b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_s_bias.sh
new file mode 100644
index 0000000000..2147bf92bd
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_s_bias.sh
@@ -0,0 +1,135 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=900
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=0.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=0
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --enable-st 0 \
+ --asr-moe 0 \
+ --asr-src 1 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
+
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_sc_moe.sh b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_sc_moe.sh
new file mode 100644
index 0000000000..f4e5854b00
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage1/cr_ctc_sc_moe.sh
@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_pre_train/merged_asr9_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=900
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=0.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=0
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 0 \
+ --enable-st 0 \
+ --asr-moe 1 \
+ --asr-src 1 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 8 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 1 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2m.sh b/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2m.sh
new file mode 100644
index 0000000000..6f9f5aee39
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2m.sh
@@ -0,0 +1,141 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_moe/merged_asr9_ast72_x_x_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_moe/merged_asr9_ast72_x_x_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=450
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=1.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=1
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+RESUME_OPT=0
+RESET_PROGRESS=1
+RESUME_CKPT="exp/europarl/best-valid-loss.pt"
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --enable-st 1 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --resume-from-checkpoint ${RESUME_CKPT} \
+ --resume-optimizer-scheduler-scaler ${RESUME_OPT} \
+ --reset-progress-stats ${RESET_PROGRESS} \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2o.sh b/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2o.sh
new file mode 100644
index 0000000000..7fb31c9a92
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage2/hent_srt_m2o.sh
@@ -0,0 +1,142 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast_de/bpe.model" # One for each language
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_no_moe/merged_asr8_ast_x_de_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_no_moe/merged_asr8_ast_x_de_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=450
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=1.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=1
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+RESUME_OPT=0
+RESET_PROGRESS=1
+RESUME_CKPT="exp/europarl/best-valid-loss.pt"
+
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --enable-st 1 \
+ --asr-moe 0 \
+ --asr-src 0 \
+ --ast-moe 0 \
+ --ast-tgt 0 \
+ --entropy-reg-asr 0.0 \
+ --entropy-reg-ast 0.0 \
+ --num-experts-asr 0 \
+ --num-experts-ast 0 \
+ --dump-moe-routing-stats 0 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --resume-from-checkpoint ${RESUME_CKPT} \
+ --resume-optimizer-scheduler-scaler ${RESUME_OPT} \
+ --reset-progress-stats ${RESET_PROGRESS} \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/train/stage2/lcma_srt.sh b/egs/europarl_st/SRT/lcma_srt/train/stage2/lcma_srt.sh
new file mode 100644
index 0000000000..90aef65721
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/train/stage2/lcma_srt.sh
@@ -0,0 +1,141 @@
+#!/usr/bin/env bash
+set -e
+export DISABLE_VERSION_CHECK=1
+
+echo "=== Training script started on $(hostname) at $(date) ==="
+log() {
+ # This function is from espnet
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%d %H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+export TRANSFORMERS_NO_GIT=1
+export GIT_DISCOVERY_ACROSS_FILESYSTEM=1
+
+ASR_BPE_MODEL="data/Europarl-ST/bpe/asr9/bpe.model"
+AST_BPE_MODEL="data/Europarl-ST/bpe/ast9/bpe.model"
+
+TRAIN_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_moe/merged_asr9_ast72_x_x_shuffled.train.jsonl.gz"
+VALID_CUTS_PATHS="data/Europarl-ST/cuts_data/asr_and_ast_moe/merged_asr9_ast72_x_x_shuffled.dev.jsonl.gz"
+
+TRAIN_PY="lcma_srt/train.py"
+
+MAX_DURATION=450
+NUM_EPOCHS=50
+BASE_LR=0.02
+START_EPOCH=1
+EXP_DIR="exp/europarl"
+mkdir -p ${EXP_DIR}
+log_path="${EXP_DIR}/run_$(date '+%Y-%m-%d_%H-%M-%S').log"
+
+MANIFEST_DIR="data/fbank"
+
+ASR_NUM_LAYERS="2,2,2,2,2"
+ASR_FF_DIM="512,768,1024,1024,1024"
+ASR_ENC_DIM="192,256,384,512,384"
+ASR_UNMASK_DIM="192,192,256,256,256"
+downsampling_factor_asr="1,2,4,8,4"
+cnn_module_kernel_asr="1,31,15,15,15"
+num_heads_asr="4,4,4,8,8"
+
+ST_NUM_LAYERS="2,2,2,2,2"
+ST_FF_DIM="512,512,256,256,256"
+ST_ENC_DIM="384,512,256,256,256"
+ST_UNMASK_DIM="256,256,256,256,192"
+downsampling_factor_st="1,2,4,4,4"
+cnn_module_kernel_st="15,31,31,15,15"
+num_heads_st="8,8,8,8,8"
+
+
+causal=0
+CHUNK_SIZE="-1"
+LEFT_CONTEXT="-1"
+
+TASK_WEIGHT_ASR=1.0
+TASK_WEIGHT_ST=1.0
+USE_TRANSDUCER=1
+USE_CTC_ASR=1
+USE_CTC_ST=1
+USE_CR_CTC=1
+USE_ATT_DEC=0
+
+RESUME_OPT=0
+RESET_PROGRESS=1
+RESUME_CKPT="exp/europarl/best-valid-loss.pt"
+
+torchrun --nproc_per_node=4 \
+ ${TRAIN_PY} \
+ --world-size 4 \
+ --num-workers 16 \
+ --num-epochs ${NUM_EPOCHS} \
+ --start-epoch ${START_EPOCH} \
+ --exp-dir ${EXP_DIR} \
+ --bpe-model-asr ${ASR_BPE_MODEL} \
+ --bpe-model-st ${AST_BPE_MODEL} \
+ --base-lr ${BASE_LR} \
+ --max-duration ${MAX_DURATION} \
+ --train-cuts-paths ${TRAIN_CUTS_PATHS} \
+ --valid-cuts-paths ${VALID_CUTS_PATHS} \
+ --utterance-min-duration 0.3 \
+ --utterance-max-duration 30.0 \
+ --manifest-dir "${MANIFEST_DIR}" \
+ --num-encoder-layers-asr ${ASR_NUM_LAYERS} \
+ --feedforward-dim-asr ${ASR_FF_DIM} \
+ --encoder-dim-asr ${ASR_ENC_DIM} \
+ --encoder-unmasked-dim-asr ${ASR_UNMASK_DIM} \
+ --num-encoder-layers-st ${ST_NUM_LAYERS} \
+ --feedforward-dim-st ${ST_FF_DIM} \
+ --encoder-dim-st ${ST_ENC_DIM} \
+ --encoder-unmasked-dim-st ${ST_UNMASK_DIM} \
+ --chunk-size "${CHUNK_SIZE}" \
+ --left-context-frames "${LEFT_CONTEXT}" \
+ --use-transducer ${USE_TRANSDUCER} \
+ --use-ctc-asr ${USE_CTC_ASR} \
+ --use-ctc-st ${USE_CTC_ST} \
+ --use-cr-ctc ${USE_CR_CTC} \
+ --use-attention-decoder ${USE_ATT_DEC} \
+ --task-weight-asr ${TASK_WEIGHT_ASR} \
+ --task-weight-st ${TASK_WEIGHT_ST} \
+ --prune-range-asr 10 \
+ --prune-range-st 10 \
+ --enable-spec-aug 0 \
+ --use-fp16 1 \
+ --causal ${causal} \
+ --full-libri 1 \
+ --ctc-loss-scale 0.1 \
+ --cr-loss-scale 0.05 \
+ --num-buckets 100 \
+ --downsampling-factor-st ${downsampling_factor_st} \
+ --cnn-module-kernel-st ${cnn_module_kernel_st} \
+ --num-heads-st ${num_heads_st} \
+ --downsampling-factor-asr ${downsampling_factor_asr} \
+ --cnn-module-kernel-asr ${cnn_module_kernel_asr} \
+ --num-heads-asr ${num_heads_asr} \
+ --freeze-asr 0 \
+ --freeze-frontend 0 \
+ --lr-epochs 6 \
+ --warm-step 2000 \
+ --output-downsampling-factor-st 1 \
+ --decoder-dim-asr 256 \
+ --decoder-dim-st 256 \
+ --joiner-dim-asr 256 \
+ --joiner-dim-st 256 \
+ --use-tgt 1 \
+ --enable-st 1 \
+ --asr-moe 1 \
+ --asr-src 1 \
+ --ast-moe 1 \
+ --ast-tgt 1 \
+ --entropy-reg-asr 0.015 \
+ --entropy-reg-ast 0.015 \
+ --num-experts-asr 8 \
+ --num-experts-ast 16 \
+ --dump-moe-routing-stats 1 \
+ --tgt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --srt-langs "en,de,es,fr,it,nl,pl,pt,ro" \
+ --resume-from-checkpoint ${RESUME_CKPT} \
+ --resume-optimizer-scheduler-scaler ${RESUME_OPT} \
+ --reset-progress-stats ${RESET_PROGRESS} \
+ > ${log_path} 2>&1
+
+echo "=== Training finished at $(date) ==="
diff --git a/egs/europarl_st/SRT/lcma_srt/zipformer.py b/egs/europarl_st/SRT/lcma_srt/zipformer.py
new file mode 100644
index 0000000000..cd12822d2a
--- /dev/null
+++ b/egs/europarl_st/SRT/lcma_srt/zipformer.py
@@ -0,0 +1,2464 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import copy
+import logging
+import math
+import random
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+from encoder_interface import EncoderInterface
+from scaling import (
+ Identity, # more friendly to backward hooks than nn.Identity(), for diagnostic reasons.
+)
+from scaling import (
+ ScaledLinear, # not as in other dirs.. just scales down initial parameter values.
+)
+from scaling import (
+ ActivationDropoutAndLinear,
+ Balancer,
+ BiasNorm,
+ ChunkCausalDepthwiseConv1d,
+ Dropout2,
+ FloatLike,
+ ScheduledFloat,
+ Whiten,
+ convert_num_channels,
+ limit_param_value,
+ penalize_abs_values_gt,
+ softmax,
+)
+from torch import Tensor, nn
+
+
+class Zipformer2(EncoderInterface):
+ """
+ Args:
+
+ Note: all "int or Tuple[int]" arguments below will be treated as lists of the same length
+ as downsampling_factor if they are single ints or one-element tuples. The length of
+ downsampling_factor defines the number of stacks.
+
+ output_downsampling_factor (int): how much to downsample at the output. Note:
+ we also downsample by a factor of 2 in the Conv2dSubsampling encoder.
+ You should probably leave this at 2.
+ downsampling_factor (Tuple[int]): downsampling factor for each encoder stack.
+ Note: this is in addition to the downsampling factor of 2 that is applied in
+ the frontend (self.encoder_embed).
+ encoder_dim (Tuple[int]): embedding dimension of each of the encoder stacks, one per
+ encoder stack.
+ num_encoder_layers (int or Tuple[int])): number of encoder layers for each stack
+ encoder_unmasked_dim (int or Tuple[int]): unmasked dimension in each of
+ the encoder stacks for purposes of per-frame dropout (recommend 256 for
+ now).
+ query_head_dim (int or Tuple[int]): dimension of query and key per attention
+ head: per stack, if a tuple..
+ pos_head_dim (int or Tuple[int]): dimension of positional-encoding projection per
+ attention head
+ value_head_dim (int or Tuple[int]): dimension of value in each attention head
+ num_heads: (int or Tuple[int]): number of heads in the self-attention mechanism.
+ Must be at least 4.
+ feedforward_dim (int or Tuple[int]): hidden dimension in feedforward modules
+ cnn_module_kernel (int or Tuple[int])): Kernel size of convolution module
+
+ pos_dim (int): the dimension of each positional-encoding vector prior to projection,
+ e.g. 128.
+
+ dropout (float): dropout rate
+ warmup_batches (float): number of batches to warm up over; this controls
+ dropout of encoder layers.
+ causal (bool): if True, support chunkwise causal convolution. This should
+ not hurt WER as no modeling power is lost, but the convolution modules will be
+ slightly slower and use more memory. Enables use of the chunk_size and
+ left_context_chunks options in forward(), which simulates streaming
+ decoding.
+ chunk_size: (list of int): only set this to other than [-1] if causal;
+ the chunk size will be randomly chosen from this list. -1 means no chunking.
+ left_context_frames: (list of int): determines the number of left-
+ context chunks for causal training; will be rounded to a number of
+ chunks. Must not be less than cnn_module_kernel (after factoring in
+ rounding and downsampling); an error will be thrown if this is violated.
+ """
+
+ def __init__(
+ self,
+ output_downsampling_factor: int = 2,
+ downsampling_factor: Tuple[int] = (2, 4),
+ encoder_dim: Union[int, Tuple[int]] = 384,
+ num_encoder_layers: Union[int, Tuple[int]] = 4,
+ encoder_unmasked_dim: Union[int, Tuple[int]] = 256,
+ query_head_dim: Union[int, Tuple[int]] = 24,
+ pos_head_dim: Union[int, Tuple[int]] = 4,
+ value_head_dim: Union[int, Tuple[int]] = 12,
+ num_heads: Union[int, Tuple[int]] = 8,
+ feedforward_dim: Union[int, Tuple[int]] = 1536,
+ cnn_module_kernel: Union[int, Tuple[int]] = 31,
+ pos_dim: int = 192,
+ dropout: FloatLike = None, # see code below for default
+ warmup_batches: float = 4000.0,
+ causal: bool = False,
+ chunk_size: Tuple[int] = [-1],
+ left_context_frames: Tuple[int] = [-1],
+ ) -> None:
+ super(Zipformer2, self).__init__()
+
+ if dropout is None:
+ dropout = ScheduledFloat((0.0, 0.3), (20000.0, 0.1))
+
+ def _to_tuple(x):
+ """Converts a single int or a 1-tuple of an int to a tuple with the same length
+ as downsampling_factor"""
+ if isinstance(x, int):
+ x = (x,)
+ if len(x) == 1:
+ x = x * len(downsampling_factor)
+ else:
+ assert len(x) == len(downsampling_factor) and isinstance(x[0], int)
+ return x
+
+ self.output_downsampling_factor = output_downsampling_factor # int
+ self.downsampling_factor = downsampling_factor # tuple
+ self.encoder_dim = encoder_dim = _to_tuple(encoder_dim) # tuple
+ self.output_dim = max(self.encoder_dim)
+ self.encoder_unmasked_dim = encoder_unmasked_dim = _to_tuple(
+ encoder_unmasked_dim
+ ) # tuple
+ num_encoder_layers = _to_tuple(num_encoder_layers)
+ self.num_encoder_layers = num_encoder_layers
+ self.query_head_dim = query_head_dim = _to_tuple(query_head_dim)
+ self.value_head_dim = value_head_dim = _to_tuple(value_head_dim)
+ pos_head_dim = _to_tuple(pos_head_dim)
+ self.num_heads = num_heads = _to_tuple(num_heads)
+ feedforward_dim = _to_tuple(feedforward_dim)
+ self.cnn_module_kernel = cnn_module_kernel = _to_tuple(cnn_module_kernel)
+
+ self.causal = causal
+ self.chunk_size = chunk_size
+ self.left_context_frames = left_context_frames
+
+ for u, d in zip(encoder_unmasked_dim, encoder_dim):
+ assert u <= d
+
+ # each one will be Zipformer2Encoder or DownsampledZipformer2Encoder
+ encoders = []
+
+ num_encoders = len(downsampling_factor)
+ for i in range(num_encoders):
+ encoder_layer = Zipformer2EncoderLayer(
+ embed_dim=encoder_dim[i],
+ pos_dim=pos_dim,
+ num_heads=num_heads[i],
+ query_head_dim=query_head_dim[i],
+ pos_head_dim=pos_head_dim[i],
+ value_head_dim=value_head_dim[i],
+ feedforward_dim=feedforward_dim[i],
+ dropout=dropout,
+ cnn_module_kernel=cnn_module_kernel[i],
+ causal=causal,
+ )
+
+ # For the segment of the warmup period, we let the Conv2dSubsampling
+ # layer learn something. Then we start to warm up the other encoders.
+ encoder = Zipformer2Encoder(
+ encoder_layer,
+ num_encoder_layers[i],
+ pos_dim=pos_dim,
+ dropout=dropout,
+ warmup_begin=warmup_batches * (i + 1) / (num_encoders + 1),
+ warmup_end=warmup_batches * (i + 2) / (num_encoders + 1),
+ final_layerdrop_rate=0.035 * (downsampling_factor[i] ** 0.5),
+ )
+
+ if downsampling_factor[i] != 1:
+ encoder = DownsampledZipformer2Encoder(
+ encoder,
+ dim=encoder_dim[i],
+ downsample=downsampling_factor[i],
+ dropout=dropout,
+ causal=causal,
+ )
+
+ encoders.append(encoder)
+
+ self.encoders = nn.ModuleList(encoders)
+
+ self.downsample_output = SimpleDownsample(
+ max(encoder_dim),
+ downsample=output_downsampling_factor,
+ dropout=dropout,
+ causal=causal,
+ )
+
+ def get_feature_masks(self, x: Tensor) -> Union[List[float], List[Tensor]]:
+ """
+ In eval mode, returns [1.0] * num_encoders; in training mode, returns a number of
+ randomized feature masks, one per encoder.
+ On e.g. 15% of frames, these masks will zero out all encoder dims larger than
+ some supplied number, e.g. >256, so in effect on those frames we are using
+ a smaller encoder dim.
+
+ We generate the random masks at this level because we want the 2 masks to 'agree'
+ all the way up the encoder stack. This will mean that the 1st mask will have
+ mask values repeated self.zipformer_subsampling_factor times.
+
+ Args:
+ x: the embeddings (needed for the shape and dtype and device), of shape
+ (1, batch_size, encoder_dims0)
+ """
+ num_encoders = len(self.encoder_dim)
+ if not self.training:
+ return [1.0] * num_encoders
+
+ (num_frames0, batch_size, _encoder_dims0) = x.shape
+
+ assert self.encoder_dim[0] == _encoder_dims0, (
+ self.encoder_dim[0],
+ _encoder_dims0,
+ )
+
+ feature_mask_dropout_prob = 0.125
+
+ # mask1 shape: (1, batch_size, 1)
+ mask1 = (
+ torch.rand(1, batch_size, 1, device=x.device) > feature_mask_dropout_prob
+ ).to(x.dtype)
+
+ # mask2 has additional sequences masked, about twice the number.
+ mask2 = torch.logical_and(
+ mask1,
+ (
+ torch.rand(1, batch_size, 1, device=x.device)
+ > feature_mask_dropout_prob
+ ).to(x.dtype),
+ )
+
+ # dim: (1, batch_size, 2)
+ mask = torch.cat((mask1, mask2), dim=-1)
+
+ feature_masks = []
+ for i in range(num_encoders):
+ channels = self.encoder_dim[i]
+ feature_mask = torch.ones(
+ 1, batch_size, channels, dtype=x.dtype, device=x.device
+ )
+ u1 = self.encoder_unmasked_dim[i]
+ u2 = u1 + (channels - u1) // 2
+
+ feature_mask[:, :, u1:u2] *= mask[..., 0:1]
+ feature_mask[:, :, u2:] *= mask[..., 1:2]
+
+ feature_masks.append(feature_mask)
+
+ return feature_masks
+
+ def get_chunk_info(self) -> Tuple[int, int]:
+ """
+ Returns chunk_size and left_context_chunks.
+ """
+ if not self.causal:
+ return -1, -1
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ assert len(self.chunk_size) == 1, self.chunk_size
+ chunk_size = self.chunk_size[0]
+ else:
+ chunk_size = random.choice(self.chunk_size)
+
+ if chunk_size == -1:
+ left_context_chunks = -1
+ else:
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ assert len(self.left_context_frames) == 1, self.left_context_frames
+ left_context_frames = self.left_context_frames[0]
+ else:
+ left_context_frames = random.choice(self.left_context_frames)
+ # Note: in Python, -1 // n == -1 for n > 0
+ left_context_chunks = left_context_frames // chunk_size
+ if left_context_chunks == 0:
+ left_context_chunks = 1
+
+ return chunk_size, left_context_chunks
+
+ def forward(
+ self,
+ x: Tensor,
+ x_lens: Tensor,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Tensor]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (seq_len, batch_size, feature_dim).
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x` before padding.
+ src_key_padding_mask:
+ The mask for padding, of shape (batch_size, seq_len); True means
+ masked position. May be None.
+ Returns:
+ Return a tuple containing 2 tensors:
+ - embeddings: its shape is (output_seq_len, batch_size, max(encoder_dim))
+ - lengths, a tensor of shape (batch_size,) containing the number
+ of frames in `embeddings` before padding.
+ """
+ outputs = []
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ feature_masks = [1.0] * len(self.encoder_dim)
+ else:
+ feature_masks = self.get_feature_masks(x)
+
+ chunk_size, left_context_chunks = self.get_chunk_info()
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ # Not support exporting a model for simulating streaming decoding
+ attn_mask = None
+ else:
+ attn_mask = self._get_attn_mask(x, chunk_size, left_context_chunks)
+
+ for i, module in enumerate(self.encoders):
+ ds = self.downsampling_factor[i]
+ x = convert_num_channels(x, self.encoder_dim[i])
+
+ x = module(
+ x,
+ chunk_size=chunk_size,
+ feature_mask=feature_masks[i],
+ src_key_padding_mask=(
+ None
+ if src_key_padding_mask is None
+ else src_key_padding_mask[..., ::ds]
+ ),
+ attn_mask=attn_mask,
+ )
+ outputs.append(x)
+
+ # if the last output has the largest dimension, x will be unchanged,
+ # it will be the same as outputs[-1]. Otherwise it will be concatenated
+ # from different pieces of 'outputs', taking each dimension from the
+ # most recent output that has it present.
+ x = self._get_full_dim_output(outputs)
+ # x = self.downsample_output(x)
+ # class Downsample has this rounding behavior..
+ # assert self.output_downsampling_factor == 2, self.output_downsampling_factor
+ # if torch.jit.is_scripting() or torch.jit.is_tracing():
+ # lengths = (x_lens + 1) // 2
+ # else:
+ # with warnings.catch_warnings():
+ # warnings.simplefilter("ignore")
+ # lengths = (x_lens + 1) // 2
+
+ # return x, lengths
+ return x, x_lens
+
+ def _get_attn_mask(
+ self, x: Tensor, chunk_size: int, left_context_chunks: int
+ ) -> Optional[Tensor]:
+ """
+ Return None if chunk_size == -1, else return attention mask of shape
+ (seq_len, seq_len), interpreted as (tgt_seq_len, src_seq_len). True
+ means a masked position.
+ Args:
+ x: embeddings after self.encoder_embed(), of shape (seq_len, batch_size, embed_dim).
+ chunk_size: chunk size, must divide
+ """
+ if chunk_size <= 0:
+ return None
+ assert all(chunk_size % d == 0 for d in self.downsampling_factor)
+ if left_context_chunks >= 0:
+ num_encoders = len(self.encoder_dim)
+ assert all(
+ chunk_size * left_context_chunks
+ >= (self.cnn_module_kernel[i] // 2) * self.downsampling_factor[i]
+ for i in range(num_encoders)
+ )
+ else:
+ left_context_chunks = 1000000
+
+ seq_len = x.shape[0]
+
+ # t is frame index, shape (seq_len,)
+ t = torch.arange(seq_len, dtype=torch.int32, device=x.device)
+ # c is chunk index for each frame, shape (seq_len,)
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ c = t // chunk_size
+ else:
+ with warnings.catch_warnings():
+ warnings.simplefilter("ignore")
+ c = t // chunk_size
+ src_c = c
+ tgt_c = c.unsqueeze(-1)
+
+ attn_mask = torch.logical_or(src_c > tgt_c, src_c < tgt_c - left_context_chunks)
+ if __name__ == "__main__":
+ logging.info(f"attn_mask = {attn_mask}")
+ return attn_mask
+
+ def _get_full_dim_output(self, outputs: List[Tensor]):
+ num_encoders = len(self.encoder_dim)
+ assert len(outputs) == num_encoders
+ output_dim = max(self.encoder_dim)
+ output_pieces = [outputs[-1]]
+ cur_dim = self.encoder_dim[-1]
+ for i in range(num_encoders - 2, -1, -1):
+ d = self.encoder_dim[i]
+ if d > cur_dim:
+ this_output = outputs[i]
+ output_pieces.append(this_output[..., cur_dim:d])
+ cur_dim = d
+ assert cur_dim == output_dim
+ return torch.cat(output_pieces, dim=-1)
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ x_lens: Tensor,
+ states: List[Tensor],
+ src_key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, Tensor, List[Tensor]]:
+ """
+ Args:
+ x:
+ The input tensor. Its shape is (seq_len, batch_size, feature_dim).
+ x_lens:
+ A tensor of shape (batch_size,) containing the number of frames in
+ `x` before padding.
+ states: list of cached tensors of all encoder layers. For layer-i,
+ states[i*6:(i+1)*6] is (cached_key, cached_nonlin_attn, cached_val1, cached_val2,
+ cached_conv1, cached_conv2).
+ src_key_padding_mask:
+ The mask for padding, of shape (batch_size, seq_len); True means
+ masked position. May be None.
+ Returns:
+ Return a tuple containing 2 tensors:
+ - embeddings: its shape is (output_seq_len, batch_size, max(encoder_dim))
+ - lengths, a tensor of shape (batch_size,) containing the number
+ of frames in `embeddings` before padding.
+ - updated states
+ """
+ outputs = []
+ new_states = []
+ layer_offset = 0
+
+ for i, module in enumerate(self.encoders):
+ num_layers = module.num_layers
+ ds = self.downsampling_factor[i]
+ x = convert_num_channels(x, self.encoder_dim[i])
+
+ x, new_layer_states = module.streaming_forward(
+ x,
+ states=states[layer_offset * 6 : (layer_offset + num_layers) * 6],
+ left_context_len=self.left_context_frames[0] // ds,
+ src_key_padding_mask=src_key_padding_mask[..., ::ds],
+ )
+ layer_offset += num_layers
+ outputs.append(x)
+ new_states += new_layer_states
+
+ # if the last output has the largest dimension, x will be unchanged,
+ # it will be the same as outputs[-1]. Otherwise it will be concatenated
+ # from different pieces of 'outputs', taking each dimension from the
+ # most recent output that has it present.
+ x = self._get_full_dim_output(outputs)
+ # x = self.downsample_output(x)
+ # class Downsample has this rounding behavior..
+ # assert self.output_downsampling_factor == 2
+ # if torch.jit.is_scripting() or torch.jit.is_tracing():
+ # lengths = (x_lens + 1) // 2
+ # else:
+ # with warnings.catch_warnings():
+ # warnings.simplefilter("ignore")
+ # lengths = (x_lens + 1) // 2
+
+ # return x, lengths, new_states
+ return x, x_lens, new_states
+
+ @torch.jit.export
+ def get_init_states(
+ self,
+ batch_size: int = 1,
+ device: torch.device = torch.device("cpu"),
+ ) -> List[Tensor]:
+ """Get initial states.
+
+ A list of cached tensors of all encoder layers. For layer-i, states[i*6:(i+1)*6]
+ is (cached_key, cached_nonlin_attn, cached_val1, cached_val2, cached_conv1, cached_conv2).
+ """
+ states = []
+ for i, module in enumerate(self.encoders):
+ num_layers = module.num_layers
+ embed_dim = self.encoder_dim[i]
+ ds = self.downsampling_factor[i]
+ num_heads = self.num_heads[i]
+ key_dim = self.query_head_dim[i] * num_heads
+ value_dim = self.value_head_dim[i] * num_heads
+ downsample_left = self.left_context_frames[0] // ds
+ nonlin_attn_head_dim = 3 * embed_dim // 4
+ conv_left_pad = self.cnn_module_kernel[i] // 2
+ for layer in range(num_layers):
+ cached_key = torch.zeros(downsample_left, batch_size, key_dim).to(
+ device
+ )
+ cached_nonlin_attn = torch.zeros(
+ 1, batch_size, downsample_left, nonlin_attn_head_dim
+ ).to(device)
+ cached_val1 = torch.zeros(downsample_left, batch_size, value_dim).to(
+ device
+ )
+ cached_val2 = torch.zeros(downsample_left, batch_size, value_dim).to(
+ device
+ )
+ cached_conv1 = torch.zeros(batch_size, embed_dim, conv_left_pad).to(
+ device
+ )
+ cached_conv2 = torch.zeros(batch_size, embed_dim, conv_left_pad).to(
+ device
+ )
+ states += [
+ cached_key,
+ cached_nonlin_attn,
+ cached_val1,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ ]
+
+ return states
+
+
+def _whitening_schedule(x: float, ratio: float = 2.0) -> ScheduledFloat:
+ return ScheduledFloat((0.0, x), (20000.0, ratio * x), default=x)
+
+
+def _balancer_schedule(min_prob: float):
+ return ScheduledFloat((0.0, 0.4), (8000.0, min_prob))
+
+
+class Zipformer2EncoderLayer(nn.Module):
+ """
+ Args:
+ embed_dim: the number of expected features in the input (required).
+ nhead: the number of heads in the multiheadattention models (required).
+ feedforward_dim: the dimension of the feedforward network model (required).
+ dropout: the dropout value (default=0.1).
+ cnn_module_kernel (int): Kernel size of convolution module (default=31).
+
+ Examples::
+ >>> encoder_layer = Zipformer2EncoderLayer(embed_dim=512, nhead=8)
+ >>> src = torch.rand(10, 32, 512)
+ >>> pos_emb = torch.rand(32, 19, 512)
+ >>> out = encoder_layer(src, pos_emb)
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ pos_dim: int,
+ num_heads: int,
+ query_head_dim: int,
+ pos_head_dim: int,
+ value_head_dim: int,
+ feedforward_dim: int,
+ dropout: FloatLike = 0.1,
+ cnn_module_kernel: int = 31,
+ causal: bool = False,
+ attention_skip_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0
+ ),
+ conv_skip_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0
+ ),
+ const_attention_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.25), (4000.0, 0.025), default=0
+ ),
+ ff2_skip_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0)
+ ),
+ ff3_skip_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0)
+ ),
+ bypass_skip_rate: FloatLike = ScheduledFloat(
+ (0.0, 0.5), (4000.0, 0.02), default=0
+ ),
+ ) -> None:
+ super(Zipformer2EncoderLayer, self).__init__()
+ self.embed_dim = embed_dim
+
+ # self.bypass implements layer skipping as well as bypass; see its default values.
+ self.bypass = BypassModule(
+ embed_dim, skip_rate=bypass_skip_rate, straight_through_rate=0
+ )
+ # bypass_mid is bypass used in the middle of the layer.
+ self.bypass_mid = BypassModule(embed_dim, straight_through_rate=0)
+
+ # skip probability for dynamic modules (meaning: anything but feedforward).
+ self.attention_skip_rate = copy.deepcopy(attention_skip_rate)
+ # an additional skip probability that applies to ConvModule to stop it from
+ # contributing too much early on.
+ self.conv_skip_rate = copy.deepcopy(conv_skip_rate)
+
+ # ff2_skip_rate is to prevent the ff2 module from having output that's too big
+ # compared to its residual.
+ self.ff2_skip_rate = copy.deepcopy(ff2_skip_rate)
+ self.ff3_skip_rate = copy.deepcopy(ff3_skip_rate)
+
+ self.const_attention_rate = copy.deepcopy(const_attention_rate)
+
+ self.self_attn_weights = RelPositionMultiheadAttentionWeights(
+ embed_dim,
+ pos_dim=pos_dim,
+ num_heads=num_heads,
+ query_head_dim=query_head_dim,
+ pos_head_dim=pos_head_dim,
+ dropout=0.0,
+ )
+
+ self.self_attn1 = SelfAttention(embed_dim, num_heads, value_head_dim)
+
+ self.self_attn2 = SelfAttention(embed_dim, num_heads, value_head_dim)
+
+ self.feed_forward1 = FeedforwardModule(
+ embed_dim, (feedforward_dim * 3) // 4, dropout
+ )
+
+ self.feed_forward2 = FeedforwardModule(embed_dim, feedforward_dim, dropout)
+
+ self.feed_forward3 = FeedforwardModule(
+ embed_dim, (feedforward_dim * 5) // 4, dropout
+ )
+
+ self.nonlin_attention = NonlinAttention(
+ embed_dim, hidden_channels=3 * embed_dim // 4
+ )
+
+ self.conv_module1 = ConvolutionModule(
+ embed_dim, cnn_module_kernel, causal=causal
+ )
+
+ self.conv_module2 = ConvolutionModule(
+ embed_dim, cnn_module_kernel, causal=causal
+ )
+
+ # TODO: remove it
+ self.bypass_scale = nn.Parameter(torch.full((embed_dim,), 0.5))
+
+ self.norm = BiasNorm(embed_dim)
+
+ self.balancer1 = Balancer(
+ embed_dim,
+ channel_dim=-1,
+ min_positive=0.45,
+ max_positive=0.55,
+ min_abs=0.2,
+ max_abs=4.0,
+ )
+
+ # balancer for output of NonlinAttentionModule
+ self.balancer_na = Balancer(
+ embed_dim,
+ channel_dim=-1,
+ min_positive=0.3,
+ max_positive=0.7,
+ min_abs=ScheduledFloat((0.0, 0.004), (4000.0, 0.02)),
+ prob=0.05, # out of concern for memory usage
+ )
+
+ # balancer for output of feedforward2, prevent it from staying too
+ # small. give this a very small probability, even at the start of
+ # training, it's to fix a rare problem and it's OK to fix it slowly.
+ self.balancer_ff2 = Balancer(
+ embed_dim,
+ channel_dim=-1,
+ min_positive=0.3,
+ max_positive=0.7,
+ min_abs=ScheduledFloat((0.0, 0.0), (4000.0, 0.1), default=0.0),
+ max_abs=2.0,
+ prob=0.05,
+ )
+
+ self.balancer_ff3 = Balancer(
+ embed_dim,
+ channel_dim=-1,
+ min_positive=0.3,
+ max_positive=0.7,
+ min_abs=ScheduledFloat((0.0, 0.0), (4000.0, 0.2), default=0.0),
+ max_abs=4.0,
+ prob=0.05,
+ )
+
+ self.whiten = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(4.0, ratio=3.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ self.balancer2 = Balancer(
+ embed_dim,
+ channel_dim=-1,
+ min_positive=0.45,
+ max_positive=0.55,
+ min_abs=0.1,
+ max_abs=4.0,
+ )
+
+ def get_sequence_dropout_mask(
+ self, x: Tensor, dropout_rate: float
+ ) -> Optional[Tensor]:
+ if (
+ dropout_rate == 0.0
+ or not self.training
+ or torch.jit.is_scripting()
+ or torch.jit.is_tracing()
+ ):
+ return None
+ batch_size = x.shape[1]
+ mask = (torch.rand(batch_size, 1, device=x.device) > dropout_rate).to(x.dtype)
+ return mask
+
+ def sequence_dropout(self, x: Tensor, dropout_rate: float) -> Tensor:
+ """
+ Apply sequence-level dropout to x.
+ x shape: (seq_len, batch_size, embed_dim)
+ """
+ dropout_mask = self.get_sequence_dropout_mask(x, dropout_rate)
+ if dropout_mask is None:
+ return x
+ else:
+ return x * dropout_mask
+
+ def forward(
+ self,
+ src: Tensor,
+ pos_emb: Tensor,
+ chunk_size: int = -1,
+ attn_mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ """
+ Pass the input through the encoder layer.
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ pos_emb: (1, 2*seq_len-1, pos_emb_dim) or (batch_size, 2*seq_len-1, pos_emb_dim)
+ chunk_size: the number of frames per chunk, of >= 0; if -1, no chunking.
+ feature_mask: something that broadcasts with src, that we'll multiply `src`
+ by at every layer: if a Tensor, likely of shape (seq_len, batch_size, embedding_dim)
+ attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) or (seq_len, seq_len),
+ interpreted as (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len).
+ True means masked position. May be None.
+ src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); True means
+ masked position. May be None.
+
+ Returns:
+ A tensor which has the same shape as src
+ """
+ src_orig = src
+
+ # dropout rate for non-feedforward submodules
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ attention_skip_rate = 0.0
+ else:
+ attention_skip_rate = (
+ float(self.attention_skip_rate) if self.training else 0.0
+ )
+
+ # attn_weights: (num_heads, batch_size, seq_len, seq_len)
+ attn_weights = self.self_attn_weights(
+ src,
+ pos_emb=pos_emb,
+ attn_mask=attn_mask,
+ key_padding_mask=src_key_padding_mask,
+ )
+
+ src = src + self.feed_forward1(src)
+
+ self_attn_dropout_mask = self.get_sequence_dropout_mask(
+ src, attention_skip_rate
+ )
+
+ selected_attn_weights = attn_weights[0:1]
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ pass
+ elif self.training and random.random() < float(self.const_attention_rate):
+ # Make attention weights constant. The intention is to
+ # encourage these modules to do something similar to an
+ # averaging-over-time operation.
+ # only need the mask, can just use the 1st one and expand later
+ selected_attn_weights = selected_attn_weights[0:1]
+ selected_attn_weights = (selected_attn_weights > 0.0).to(
+ selected_attn_weights.dtype
+ )
+ selected_attn_weights = selected_attn_weights * (
+ 1.0 / selected_attn_weights.sum(dim=-1, keepdim=True)
+ )
+
+ na = self.balancer_na(self.nonlin_attention(src, selected_attn_weights))
+
+ src = src + (
+ na if self_attn_dropout_mask is None else na * self_attn_dropout_mask
+ )
+
+ self_attn = self.self_attn1(src, attn_weights)
+
+ src = src + (
+ self_attn
+ if self_attn_dropout_mask is None
+ else self_attn * self_attn_dropout_mask
+ )
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ conv_skip_rate = 0.0
+ else:
+ conv_skip_rate = float(self.conv_skip_rate) if self.training else 0.0
+ src = src + self.sequence_dropout(
+ self.conv_module1(
+ src, chunk_size=chunk_size, src_key_padding_mask=src_key_padding_mask
+ ),
+ conv_skip_rate,
+ )
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ ff2_skip_rate = 0.0
+ else:
+ ff2_skip_rate = float(self.ff2_skip_rate) if self.training else 0.0
+ src = src + self.sequence_dropout(
+ self.balancer_ff2(self.feed_forward2(src)), ff2_skip_rate
+ )
+
+ # bypass in the middle of the layer.
+ src = self.bypass_mid(src_orig, src)
+
+ self_attn = self.self_attn2(src, attn_weights)
+
+ src = src + (
+ self_attn
+ if self_attn_dropout_mask is None
+ else self_attn * self_attn_dropout_mask
+ )
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ conv_skip_rate = 0.0
+ else:
+ conv_skip_rate = float(self.conv_skip_rate) if self.training else 0.0
+ src = src + self.sequence_dropout(
+ self.conv_module2(
+ src, chunk_size=chunk_size, src_key_padding_mask=src_key_padding_mask
+ ),
+ conv_skip_rate,
+ )
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ ff3_skip_rate = 0.0
+ else:
+ ff3_skip_rate = float(self.ff3_skip_rate) if self.training else 0.0
+ src = src + self.sequence_dropout(
+ self.balancer_ff3(self.feed_forward3(src)), ff3_skip_rate
+ )
+
+ src = self.balancer1(src)
+ src = self.norm(src)
+
+ src = self.bypass(src_orig, src)
+
+ src = self.balancer2(src)
+ src = self.whiten(src)
+
+ return src
+
+ def streaming_forward(
+ self,
+ src: Tensor,
+ pos_emb: Tensor,
+ cached_key: Tensor,
+ cached_nonlin_attn: Tensor,
+ cached_val1: Tensor,
+ cached_val2: Tensor,
+ cached_conv1: Tensor,
+ cached_conv2: Tensor,
+ left_context_len: int,
+ src_key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ """Pass the input through the encoder layer in streaming forward mode.
+
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ pos_emb: (1, left_context_len+2*seq_len-1, pos_emb_dim) or
+ (batch_size, left_context_len+2*seq_len-1, pos_emb_dim)
+ cached_key: cached attention key tensor of left context,
+ of shape (left_context_len, batch_size, key_dim)
+ cached_nonlin_attn: left context for nonlin_attention module, a Tensor of shape
+ (num_heads, batch_size, left_context_len, head_dim)
+ cached_val1: cached left context for the first attention module,
+ of shape (left_context_len, batch_size, value_dim)
+ cached_val2: cached left context for the second attention module,
+ of shape (left_context_len, batch_size, value_dim)
+ cached_conv1: cached left context for the first convolution module,
+ of shape (batch_size, channels, left_pad)
+ cached_conv2: cached left context for the second convolution module,
+ of shape (batch_size, channels, left_pad)
+ left_context_len: number of left context frames.
+ src_key_padding_mask: the mask for padding, of shape
+ (batch_size, left_context_len + seq_len); True means masked position.
+ May be None.
+
+ Returns:
+ - x, with the same shape as src
+ - updated cached_key
+ - updated cached_nonlin_attn
+ - updated cached_val1
+ - updated cached_val2
+ - updated cached_conv1
+ - updated cached_conv2
+ """
+ src_orig = src
+
+ # attn_weights: (num_heads, batch_size, seq_len, seq_len)
+ attn_weights, cached_key = self.self_attn_weights.streaming_forward(
+ src,
+ pos_emb=pos_emb,
+ cached_key=cached_key,
+ left_context_len=left_context_len,
+ key_padding_mask=src_key_padding_mask,
+ )
+
+ src = src + self.feed_forward1(src)
+
+ na, cached_nonlin_attn = self.nonlin_attention.streaming_forward(
+ src,
+ attn_weights[0:1],
+ cached_x=cached_nonlin_attn,
+ left_context_len=left_context_len,
+ )
+ src = src + na
+
+ self_attn, cached_val1 = self.self_attn1.streaming_forward(
+ src,
+ attn_weights=attn_weights,
+ cached_val=cached_val1,
+ left_context_len=left_context_len,
+ )
+ src = src + self_attn
+
+ src_conv, cached_conv1 = self.conv_module1.streaming_forward(
+ src,
+ cache=cached_conv1,
+ src_key_padding_mask=src_key_padding_mask[:, left_context_len:],
+ )
+ src = src + src_conv
+
+ src = src + self.feed_forward2(src)
+
+ # bypass in the middle of the layer.
+ src = self.bypass_mid(src_orig, src)
+
+ self_attn, cached_val2 = self.self_attn2.streaming_forward(
+ src,
+ attn_weights=attn_weights,
+ cached_val=cached_val2,
+ left_context_len=left_context_len,
+ )
+ src = src + self_attn
+
+ src_conv, cached_conv2 = self.conv_module2.streaming_forward(
+ src,
+ cache=cached_conv2,
+ src_key_padding_mask=src_key_padding_mask[:, left_context_len:],
+ )
+ src = src + src_conv
+
+ src = src + self.feed_forward3(src)
+
+ src = self.norm(src)
+
+ src = self.bypass(src_orig, src)
+
+ return (
+ src,
+ cached_key,
+ cached_nonlin_attn,
+ cached_val1,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ )
+
+
+class Zipformer2Encoder(nn.Module):
+ r"""Zipformer2Encoder is a stack of N encoder layers
+
+ Args:
+ encoder_layer: an instance of the Zipformer2EncoderLayer() class (required).
+ num_layers: the number of sub-encoder-layers in the encoder (required).
+ pos_dim: the dimension for the relative positional encoding
+
+ Examples::
+ >>> encoder_layer = Zipformer2EncoderLayer(embed_dim=512, nhead=8)
+ >>> zipformer_encoder = Zipformer2Encoder(encoder_layer, num_layers=6)
+ >>> src = torch.rand(10, 32, 512)
+ >>> out = zipformer_encoder(src)
+ """
+
+ def __init__(
+ self,
+ encoder_layer: nn.Module,
+ num_layers: int,
+ pos_dim: int,
+ dropout: float,
+ warmup_begin: float,
+ warmup_end: float,
+ initial_layerdrop_rate: float = 0.5,
+ final_layerdrop_rate: float = 0.05,
+ ) -> None:
+ super().__init__()
+ self.encoder_pos = CompactRelPositionalEncoding(
+ pos_dim, dropout_rate=0.15, length_factor=1.0
+ )
+
+ self.layers = nn.ModuleList(
+ [copy.deepcopy(encoder_layer) for i in range(num_layers)]
+ )
+ self.num_layers = num_layers
+
+ assert 0 <= warmup_begin <= warmup_end, (warmup_begin, warmup_end)
+
+ delta = (1.0 / num_layers) * (warmup_end - warmup_begin)
+ cur_begin = warmup_begin # interpreted as a training batch index
+ for i in range(num_layers):
+ cur_end = cur_begin + delta
+ self.layers[i].bypass.skip_rate = ScheduledFloat(
+ (cur_begin, initial_layerdrop_rate),
+ (cur_end, final_layerdrop_rate),
+ default=0.0,
+ )
+ cur_begin = cur_end
+
+ def forward(
+ self,
+ src: Tensor,
+ chunk_size: int = -1,
+ feature_mask: Union[Tensor, float] = 1.0,
+ attn_mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ r"""Pass the input through the encoder layers in turn.
+
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ chunk_size: the number of frames per chunk, of >= 0; if -1, no chunking.
+ feature_mask: something that broadcasts with src, that we'll multiply `src`
+ by at every layer: if a Tensor, likely of shape (seq_len, batch_size, embedding_dim)
+ attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) or (seq_len, seq_len),
+ interpreted as (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len).
+ True means masked position. May be None.
+ src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); True means
+ masked position. May be None.
+
+ Returns: a Tensor with the same shape as src.
+ """
+ pos_emb = self.encoder_pos(src)
+ output = src
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ output = output * feature_mask
+
+ for i, mod in enumerate(self.layers):
+ output = mod(
+ output,
+ pos_emb,
+ chunk_size=chunk_size,
+ attn_mask=attn_mask,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+
+ if not torch.jit.is_scripting() and not torch.jit.is_tracing():
+ output = output * feature_mask
+
+ return output
+
+ def streaming_forward(
+ self,
+ src: Tensor,
+ states: List[Tensor],
+ left_context_len: int,
+ src_key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, List[Tensor]]:
+ r"""Pass the input through the encoder layers in turn.
+
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ states: list of cached tensors of N encoder layers. For layer-i, states[i*6:(i+1)*6] is
+ (cached_key, cached_nonlin_attn, cached_val1, cached_val2, cached_conv1, cached_conv2).
+ left_context_len: Number of left context frames.
+ src_key_padding_mask: the mask for padding, of shape
+ (batch_size, left_context_len + seq_len); True means masked position.
+ May be None.
+
+ Returns:
+ - output, a Tensor with the same shape as src.
+ - updated states
+ """
+ pos_emb = self.encoder_pos(src, left_context_len)
+ output = src
+
+ new_states = []
+ for i, mod in enumerate(self.layers):
+ (
+ cached_key,
+ cached_nonlin_attn,
+ cached_val1,
+ cached_val2,
+ cached_conv1,
+ cached_conv2,
+ ) = states[i * 6 : (i + 1) * 6]
+ (
+ output,
+ new_cached_key,
+ new_cached_nonlin_attn,
+ new_cached_val1,
+ new_cached_val2,
+ new_cached_conv1,
+ new_cached_conv2,
+ ) = mod.streaming_forward(
+ output,
+ pos_emb,
+ cached_key=cached_key,
+ cached_nonlin_attn=cached_nonlin_attn,
+ cached_val1=cached_val1,
+ cached_val2=cached_val2,
+ cached_conv1=cached_conv1,
+ cached_conv2=cached_conv2,
+ left_context_len=left_context_len,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+ new_states += [
+ new_cached_key,
+ new_cached_nonlin_attn,
+ new_cached_val1,
+ new_cached_val2,
+ new_cached_conv1,
+ new_cached_conv2,
+ ]
+
+ return output, new_states
+
+
+class BypassModule(nn.Module):
+ """
+ An nn.Module that implements a learnable bypass scale, and also randomized per-sequence
+ layer-skipping. The bypass is limited during early stages of training to be close to
+ "straight-through", i.e. to not do the bypass operation much initially, in order to
+ force all the modules to learn something.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ skip_rate: FloatLike = 0.0,
+ straight_through_rate: FloatLike = 0.0,
+ scale_min: FloatLike = ScheduledFloat((0.0, 0.9), (20000.0, 0.2), default=0),
+ scale_max: FloatLike = 1.0,
+ ):
+ super().__init__()
+ self.bypass_scale = nn.Parameter(torch.full((embed_dim,), 0.5))
+ self.skip_rate = copy.deepcopy(skip_rate)
+ self.straight_through_rate = copy.deepcopy(straight_through_rate)
+ self.scale_min = copy.deepcopy(scale_min)
+ self.scale_max = copy.deepcopy(scale_max)
+
+ def _get_bypass_scale(self, batch_size: int):
+ # returns bypass-scale of shape (num_channels,),
+ # or (batch_size, num_channels,). This is actually the
+ # scale on the non-residual term, so 0 corresponds to bypassing
+ # this module.
+ if torch.jit.is_scripting() or torch.jit.is_tracing() or not self.training:
+ return self.bypass_scale
+ else:
+ ans = limit_param_value(
+ self.bypass_scale, min=float(self.scale_min), max=float(self.scale_max)
+ )
+ skip_rate = float(self.skip_rate)
+ if skip_rate != 0.0:
+ mask = torch.rand((batch_size, 1), device=ans.device) > skip_rate
+ ans = ans * mask
+ # now ans is of shape (batch_size, num_channels), and is zero for sequences
+ # on which we have randomly chosen to do layer-skipping.
+ straight_through_rate = float(self.straight_through_rate)
+ if straight_through_rate != 0.0:
+ mask = (
+ torch.rand((batch_size, 1), device=ans.device)
+ < straight_through_rate
+ )
+ ans = torch.maximum(ans, mask.to(ans.dtype))
+ return ans
+
+ def forward(self, src_orig: Tensor, src: Tensor):
+ """
+ Args: src_orig and src are both of shape (seq_len, batch_size, num_channels)
+ Returns: something with the same shape as src and src_orig
+ """
+ bypass_scale = self._get_bypass_scale(src.shape[1])
+ return src_orig + (src - src_orig) * bypass_scale
+
+
+class DownsampledZipformer2Encoder(nn.Module):
+ r"""
+ DownsampledZipformer2Encoder is a zipformer encoder evaluated at a reduced frame rate,
+ after convolutional downsampling, and then upsampled again at the output, and combined
+ with the origin input, so that the output has the same shape as the input.
+ """
+
+ def __init__(
+ self,
+ encoder: nn.Module,
+ dim: int,
+ downsample: int,
+ dropout: FloatLike,
+ causal: bool,
+ ):
+ super(DownsampledZipformer2Encoder, self).__init__()
+ self.downsample_factor = downsample
+ self.downsample = SimpleDownsample(dim, downsample, dropout, causal)
+ self.num_layers = encoder.num_layers
+ self.encoder = encoder
+ self.upsample = SimpleUpsample(dim, downsample)
+ self.out_combiner = BypassModule(dim, straight_through_rate=0)
+
+ def forward(
+ self,
+ src: Tensor,
+ chunk_size: int = -1,
+ feature_mask: Union[Tensor, float] = 1.0,
+ attn_mask: Optional[Tensor] = None,
+ src_key_padding_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ r"""Downsample, go through encoder, upsample.
+
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ feature_mask: something that broadcasts with src, that we'll multiply `src`
+ by at every layer: if a Tensor, likely of shape (seq_len, batch_size, embedding_dim)
+ attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) or (seq_len, seq_len),
+ interpreted as (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len).
+ True means masked position. May be None.
+ src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); True means
+ masked position. May be None.
+
+ Returns: a Tensor with the same shape as src.
+ """
+ src_orig = src
+ src = self.downsample(src)
+ ds = self.downsample_factor
+ if attn_mask is not None:
+ attn_mask = attn_mask[::ds, ::ds]
+
+ src = self.encoder(
+ src,
+ chunk_size=chunk_size // ds,
+ feature_mask=feature_mask,
+ attn_mask=attn_mask,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+ src = self.upsample(src)
+ # remove any extra frames that are not a multiple of downsample_factor
+ src = src[: src_orig.shape[0]]
+
+ return self.out_combiner(src_orig, src)
+
+ def streaming_forward(
+ self,
+ src: Tensor,
+ states: List[Tensor],
+ left_context_len: int,
+ src_key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, List[Tensor]]:
+ r"""Downsample, go through encoder, upsample, in streaming forward mode.
+
+ Args:
+ src: the sequence to the encoder (required): shape (seq_len, batch_size, embedding_dim).
+ states: list of cached tensors of N encoder layers. For layer-i, states[i*6:(i+1)*6] is
+ (cached_key, cached_nonlin_attn, cached_val1, cached_val2, cached_conv1, cached_conv2).
+ left_context_len: Number of left context frames.
+ src_key_padding_mask: the mask for padding, of shape (batch_size, left_context_len+seq_len);
+ True means masked position. May be None.
+
+ Returns:
+ - output, a Tensor with the same shape as src.
+ - updated states
+ """
+ src_orig = src
+ src = self.downsample(src)
+
+ src, new_states = self.encoder.streaming_forward(
+ src,
+ states=states,
+ left_context_len=left_context_len,
+ src_key_padding_mask=src_key_padding_mask,
+ )
+ src = self.upsample(src)
+ # remove any extra frames that are not a multiple of downsample_factor
+ src = src[: src_orig.shape[0]]
+
+ return self.out_combiner(src_orig, src), new_states
+
+
+class SimpleDownsample(torch.nn.Module):
+ """
+ Does downsampling with attention, by weighted sum, and a projection..
+ """
+
+ def __init__(
+ self, channels: int, downsample: int, dropout: FloatLike, causal: bool
+ ):
+ super(SimpleDownsample, self).__init__()
+
+ self.causal = causal
+ self.bias = nn.Parameter(torch.zeros(downsample))
+
+ self.name = None # will be set from training code
+ self.dropout = copy.deepcopy(dropout)
+
+ self.downsample = downsample
+
+ def forward(self, src: Tensor) -> Tensor:
+ """
+ x: (seq_len, batch_size, in_channels)
+ Returns a tensor of shape
+ ( (seq_len+downsample-1)//downsample, batch_size, channels)
+ """
+ (seq_len, batch_size, in_channels) = src.shape
+ ds = self.downsample
+ d_seq_len = (seq_len + ds - 1) // ds
+
+ # Pad to an exact multiple of self.downsample
+ # right-pad src, repeating the last element.
+ pad = d_seq_len * ds - seq_len
+
+ if self.causal and torch.jit.is_tracing():
+ assert (
+ pad == 0
+ ), f"pad should be zero for exporting streaming models. Given {pad}"
+
+ # If we are exporting a streaming model, then we skip the if statement
+ if not self.causal or not torch.jit.is_tracing():
+ src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2])
+ src = torch.cat((src, src_extra), dim=0)
+
+ assert src.shape[0] == d_seq_len * ds, (src.shape, d_seq_len, ds)
+
+ src = src.reshape(d_seq_len, ds, batch_size, in_channels)
+
+ weights = self.bias.softmax(dim=0)
+ # weights: (downsample, 1, 1)
+ weights = weights.unsqueeze(-1).unsqueeze(-1)
+
+ # ans1 is the first `in_channels` channels of the output
+ ans = (src * weights).sum(dim=1)
+
+ return ans
+
+
+class SimpleUpsample(torch.nn.Module):
+ """
+ A very simple form of upsampling that mostly just repeats the input, but
+ also adds a position-specific bias.
+ """
+
+ def __init__(self, num_channels: int, upsample: int):
+ super(SimpleUpsample, self).__init__()
+ self.upsample = upsample
+
+ def forward(self, src: Tensor) -> Tensor:
+ """
+ x: (seq_len, batch_size, num_channels)
+ Returns a tensor of shape
+ ( (seq_len*upsample), batch_size, num_channels)
+ """
+ upsample = self.upsample
+ (seq_len, batch_size, num_channels) = src.shape
+ src = src.unsqueeze(1).expand(seq_len, upsample, batch_size, num_channels)
+ src = src.reshape(seq_len * upsample, batch_size, num_channels)
+ return src
+
+
+class CompactRelPositionalEncoding(torch.nn.Module):
+ """
+ Relative positional encoding module. This version is "compact" meaning it is able to encode
+ the important information about the relative position in a relatively small number of dimensions.
+ The goal is to make it so that small differences between large relative offsets (e.g. 1000 vs. 1001)
+ make very little difference to the embedding. Such differences were potentially important
+ when encoding absolute position, but not important when encoding relative position because there
+ is now no need to compare two large offsets with each other.
+
+ Our embedding works by projecting the interval [-infinity,infinity] to a finite interval
+ using the atan() function, before doing the Fourier transform of that fixed interval. The
+ atan() function would compress the "long tails" too small,
+ making it hard to distinguish between different magnitudes of large offsets, so we use a logarithmic
+ function to compress large offsets to a smaller range before applying atan().
+ Scalings are chosen in such a way that the embedding can clearly distinguish individual offsets as long
+ as they are quite close to the origin, e.g. abs(offset) <= about sqrt(embedding_dim)
+
+
+ Args:
+ embed_dim: Embedding dimension.
+ dropout_rate: Dropout rate.
+ max_len: Maximum input length: just a heuristic for initialization.
+ length_factor: a heuristic scale (should be >= 1.0) which, if larger, gives
+ less weight to small differences of offset near the origin.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ dropout_rate: FloatLike,
+ max_len: int = 1000,
+ length_factor: float = 1.0,
+ ) -> None:
+ """Construct a CompactRelPositionalEncoding object."""
+ super(CompactRelPositionalEncoding, self).__init__()
+ self.embed_dim = embed_dim
+ assert embed_dim % 2 == 0, embed_dim
+ self.dropout = Dropout2(dropout_rate)
+ self.pe = None
+ assert length_factor >= 1.0, length_factor
+ self.length_factor = length_factor
+ self.extend_pe(torch.tensor(0.0).expand(max_len))
+
+ def extend_pe(self, x: Tensor, left_context_len: int = 0) -> None:
+ """Reset the positional encodings."""
+ T = x.size(0) + left_context_len
+
+ if self.pe is not None:
+ # self.pe contains both positive and negative parts
+ # the length of self.pe is 2 * input_len - 1
+ if self.pe.size(0) >= T * 2 - 1:
+ self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+ return
+
+ # if T == 4, x would contain [ -3, -2, 1, 0, 1, 2, 3 ]
+ x = torch.arange(-(T - 1), T, device=x.device).to(torch.float32).unsqueeze(1)
+
+ freqs = 1 + torch.arange(self.embed_dim // 2, device=x.device)
+
+ # `compression_length` this is arbitrary/heuristic, if it is larger we have more resolution
+ # for small time offsets but less resolution for large time offsets.
+ compression_length = self.embed_dim**0.5
+ # x_compressed, like X, goes from -infinity to infinity as T goes from -infinity to infinity;
+ # but it does so more slowly than T for large absolute values of T.
+ # The formula is chosen so that d(x_compressed )/dx is 1 around x == 0, which
+ # is important.
+ x_compressed = (
+ compression_length
+ * x.sign()
+ * ((x.abs() + compression_length).log() - math.log(compression_length))
+ )
+
+ # if self.length_factor == 1.0, then length_scale is chosen so that the
+ # FFT can exactly separate points close to the origin (T == 0). So this
+ # part of the formulation is not really heuristic.
+ # But empirically, for ASR at least, length_factor > 1.0 seems to work better.
+ length_scale = self.length_factor * self.embed_dim / (2.0 * math.pi)
+
+ # note for machine implementations: if atan is not available, we can use:
+ # x.sign() * ((1 / (x.abs() + 1)) - 1) * (-math.pi/2)
+ # check on wolframalpha.com: plot(sign(x) * (1 / ( abs(x) + 1) - 1 ) * -pi/2 , atan(x))
+ x_atan = (x_compressed / length_scale).atan() # results between -pi and pi
+
+ cosines = (x_atan * freqs).cos()
+ sines = (x_atan * freqs).sin()
+
+ pe = torch.zeros(x.shape[0], self.embed_dim, device=x.device)
+ pe[:, 0::2] = cosines
+ pe[:, 1::2] = sines
+ pe[:, -1] = 1.0 # for bias.
+
+ self.pe = pe.to(dtype=x.dtype)
+
+ def forward(self, x: Tensor, left_context_len: int = 0) -> Tensor:
+ """Create positional encoding.
+
+ Args:
+ x (Tensor): Input tensor (time, batch, `*`).
+ left_context_len: (int): Length of cached left context.
+
+ Returns:
+ positional embedding, of shape (batch, left_context_len + 2*time-1, `*`).
+ """
+ self.extend_pe(x, left_context_len)
+ x_size_left = x.size(0) + left_context_len
+ # length of positive side: x.size(0) + left_context_len
+ # length of negative side: x.size(0)
+ pos_emb = self.pe[
+ self.pe.size(0) // 2
+ - x_size_left
+ + 1 : self.pe.size(0) // 2 # noqa E203
+ + x.size(0),
+ :,
+ ]
+ pos_emb = pos_emb.unsqueeze(0)
+ return self.dropout(pos_emb)
+
+
+class RelPositionMultiheadAttentionWeights(nn.Module):
+ r"""Module that computes multi-head attention weights with relative position encoding.
+ Various other modules consume the resulting attention weights: see, for example, the
+ SimpleAttention module which allows you to compute conventional attention.
+
+ This is a quite heavily modified from: "Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context",
+ we have to write up the differences.
+
+
+ Args:
+ embed_dim: number of channels at the input to this module, e.g. 256
+ pos_dim: dimension of the positional encoding vectors, e.g. 128.
+ num_heads: number of heads to compute weights for, e.g. 8
+ query_head_dim: dimension of the query (and key), per head. e.g. 24.
+ pos_head_dim: dimension of the projected positional encoding per head, e.g. 4.
+ dropout: dropout probability for attn_output_weights. Default: 0.0.
+ pos_emb_skip_rate: probability for skipping the pos_emb part of the scores on
+ any given call to forward(), in training time.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ pos_dim: int,
+ num_heads: int,
+ query_head_dim: int,
+ pos_head_dim: int,
+ dropout: float = 0.0,
+ pos_emb_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.0)),
+ ) -> None:
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.query_head_dim = query_head_dim
+ self.pos_head_dim = pos_head_dim
+ self.dropout = dropout
+ self.pos_emb_skip_rate = copy.deepcopy(pos_emb_skip_rate)
+ self.name = None # will be overwritten in training code; for diagnostics.
+
+ key_head_dim = query_head_dim
+ in_proj_dim = (query_head_dim + key_head_dim + pos_head_dim) * num_heads
+
+ # the initial_scale is supposed to take over the "scaling" factor of
+ # head_dim ** -0.5 that has been used in previous forms of attention,
+ # dividing it between the query and key. Note: this module is intended
+ # to be used with the ScaledAdam optimizer; with most other optimizers,
+ # it would be necessary to apply the scaling factor in the forward function.
+ self.in_proj = ScaledLinear(
+ embed_dim, in_proj_dim, bias=True, initial_scale=query_head_dim**-0.25
+ )
+
+ self.whiten_keys = Whiten(
+ num_groups=num_heads,
+ whitening_limit=_whitening_schedule(3.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.025,
+ )
+
+ # add a balancer for the keys that runs with very small probability, and
+ # tries to enforce that all dimensions have mean around zero. The
+ # weights produced by this module are invariant to adding a constant to
+ # the keys, so the derivative of the bias is mathematically zero; but
+ # due to how Adam/ScaledAdam work, it can learn a fairly large nonzero
+ # bias because the small numerical roundoff tends to have a non-random
+ # sign. This module is intended to prevent that. Use a very small
+ # probability; that should be sufficient to fix the problem.
+ self.balance_keys = Balancer(
+ key_head_dim * num_heads,
+ channel_dim=-1,
+ min_positive=0.4,
+ max_positive=0.6,
+ min_abs=0.0,
+ max_abs=100.0,
+ prob=0.025,
+ )
+
+ # linear transformation for positional encoding.
+ self.linear_pos = ScaledLinear(
+ pos_dim, num_heads * pos_head_dim, bias=False, initial_scale=0.05
+ )
+
+ # the following are for diagnostics only, see --print-diagnostics option
+ self.copy_pos_query = Identity()
+ self.copy_query = Identity()
+
+ def forward(
+ self,
+ x: Tensor,
+ pos_emb: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ attn_mask: Optional[Tensor] = None,
+ ) -> Tensor:
+ r"""
+ Args:
+ x: input of shape (seq_len, batch_size, embed_dim)
+ pos_emb: Positional embedding tensor, of shape (1, 2*seq_len - 1, pos_dim)
+ key_padding_mask: a bool tensor of shape (batch_size, seq_len). Positions that
+ are True in this mask will be ignored as sources in the attention weighting.
+ attn_mask: mask of shape (seq_len, seq_len) or (batch_size, seq_len, seq_len),
+ interpreted as ([batch_size,] tgt_seq_len, src_seq_len)
+ saying which positions are allowed to attend to which other positions.
+ Returns:
+ a tensor of attention weights, of shape (hum_heads, batch_size, seq_len, seq_len)
+ interpreted as (hum_heads, batch_size, tgt_seq_len, src_seq_len).
+ """
+ x = self.in_proj(x)
+ query_head_dim = self.query_head_dim
+ pos_head_dim = self.pos_head_dim
+ num_heads = self.num_heads
+
+ seq_len, batch_size, _ = x.shape
+
+ query_dim = query_head_dim * num_heads
+
+ # self-attention
+ q = x[..., 0:query_dim]
+ k = x[..., query_dim : 2 * query_dim]
+ # p is the position-encoding query
+ p = x[..., 2 * query_dim :]
+ assert p.shape[-1] == num_heads * pos_head_dim, (
+ p.shape[-1],
+ num_heads,
+ pos_head_dim,
+ )
+
+ q = self.copy_query(q) # for diagnostics only, does nothing.
+ k = self.whiten_keys(self.balance_keys(k)) # does nothing in the forward pass.
+ p = self.copy_pos_query(p) # for diagnostics only, does nothing.
+
+ q = q.reshape(seq_len, batch_size, num_heads, query_head_dim)
+ p = p.reshape(seq_len, batch_size, num_heads, pos_head_dim)
+ k = k.reshape(seq_len, batch_size, num_heads, query_head_dim)
+
+ # time1 refers to target, time2 refers to source.
+ q = q.permute(2, 1, 0, 3) # (head, batch, time1, query_head_dim)
+ p = p.permute(2, 1, 0, 3) # (head, batch, time1, pos_head_dim)
+ k = k.permute(2, 1, 3, 0) # (head, batch, d_k, time2)
+
+ attn_scores = torch.matmul(q, k)
+
+ use_pos_scores = False
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ # We can't put random.random() in the same line
+ use_pos_scores = True
+ elif not self.training or random.random() >= float(self.pos_emb_skip_rate):
+ use_pos_scores = True
+
+ if use_pos_scores:
+ pos_emb = self.linear_pos(pos_emb)
+ seq_len2 = 2 * seq_len - 1
+ pos_emb = pos_emb.reshape(-1, seq_len2, num_heads, pos_head_dim).permute(
+ 2, 0, 3, 1
+ )
+ # pos shape now: (head, {1 or batch_size}, pos_dim, seq_len2)
+
+ # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, batch, time1, seq_len2)
+ # [where seq_len2 represents relative position.]
+ pos_scores = torch.matmul(p, pos_emb)
+ # the following .as_strided() expression converts the last axis of pos_scores from relative
+ # to absolute position. I don't know whether I might have got the time-offsets backwards or
+ # not, but let this code define which way round it is supposed to be.
+ if torch.jit.is_tracing():
+ (num_heads, batch_size, time1, n) = pos_scores.shape
+ rows = torch.arange(start=time1 - 1, end=-1, step=-1)
+ cols = torch.arange(seq_len)
+ rows = rows.repeat(batch_size * num_heads).unsqueeze(-1)
+ indexes = rows + cols
+ pos_scores = pos_scores.reshape(-1, n)
+ pos_scores = torch.gather(pos_scores, dim=1, index=indexes)
+ pos_scores = pos_scores.reshape(num_heads, batch_size, time1, seq_len)
+ else:
+ pos_scores = pos_scores.as_strided(
+ (num_heads, batch_size, seq_len, seq_len),
+ (
+ pos_scores.stride(0),
+ pos_scores.stride(1),
+ pos_scores.stride(2) - pos_scores.stride(3),
+ pos_scores.stride(3),
+ ),
+ storage_offset=pos_scores.stride(3) * (seq_len - 1),
+ )
+
+ attn_scores = attn_scores + pos_scores
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ pass
+ elif self.training and random.random() < 0.1:
+ # This is a harder way of limiting the attention scores to not be
+ # too large. It incurs a penalty if any of them has an absolute
+ # value greater than 50.0. this should be outside the normal range
+ # of the attention scores. We use this mechanism instead of, say,
+ # something added to the loss function involving the entropy,
+ # because once the entropy gets very small gradients through the
+ # softmax can become very small, and we'd get zero derivatives. The
+ # choices of 1.0e-04 as the scale on the penalty makes this
+ # mechanism vulnerable to the absolute scale of the loss function,
+ # but we view this as a failsafe to avoid "implausible" parameter
+ # values rather than a regularization method that should be active
+ # under normal circumstances.
+ attn_scores = penalize_abs_values_gt(
+ attn_scores, limit=25.0, penalty=1.0e-04, name=self.name
+ )
+
+ assert attn_scores.shape == (num_heads, batch_size, seq_len, seq_len)
+
+ if attn_mask is not None:
+ assert attn_mask.dtype == torch.bool
+ # use -1000 to avoid nan's where attn_mask and key_padding_mask make
+ # all scores zero. It's important that this be large enough that exp(-1000)
+ # is exactly zero, for reasons related to const_attention_rate, it
+ # compares the final weights with zero.
+ attn_scores = attn_scores.masked_fill(attn_mask, -1000)
+
+ if key_padding_mask is not None:
+ assert key_padding_mask.shape == (
+ batch_size,
+ seq_len,
+ ), key_padding_mask.shape
+ attn_scores = attn_scores.masked_fill(
+ key_padding_mask.unsqueeze(1),
+ -1000,
+ )
+
+ # We use our own version of softmax, defined in scaling.py, which should
+ # save a little of the memory used in backprop by, if we are in
+ # automatic mixed precision mode (amp / autocast), by only storing the
+ # half-precision output for backprop purposes.
+ attn_weights = softmax(attn_scores, dim=-1)
+
+ if torch.jit.is_scripting() or torch.jit.is_tracing():
+ pass
+ elif random.random() < 0.001 and not self.training:
+ self._print_attn_entropy(attn_weights)
+
+ attn_weights = nn.functional.dropout(
+ attn_weights, p=self.dropout, training=self.training
+ )
+
+ return attn_weights
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ pos_emb: Tensor,
+ cached_key: Tensor,
+ left_context_len: int,
+ key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ r"""
+ Args:
+ x: input of shape (seq_len, batch_size, embed_dim)
+ pos_emb: Positional embedding tensor, of shape (1, left_context_len+2*seq_len-1, pos_dim)
+ cached_key: cached attention key tensor of left context,
+ of shape (left_context_len, batch_size, key_dim)
+ left_context_len: number of left context frames.
+ key_padding_mask: a bool tensor of shape (batch_size, seq_len). Positions that
+ are True in this mask will be ignored as sources in the attention weighting.
+
+ Returns:
+ - attention weights, of shape (hum_heads, batch_size, seq_len, seq_len2),
+ interpreted as (hum_heads, batch_size, tgt_seq_len, src_seq_len).
+ - updated cached attention key tensor of left context.
+ """
+ x = self.in_proj(x)
+ query_head_dim = self.query_head_dim
+ pos_head_dim = self.pos_head_dim
+ num_heads = self.num_heads
+
+ seq_len, batch_size, _ = x.shape
+
+ query_dim = query_head_dim * num_heads
+
+ # self-attention
+ q = x[..., 0:query_dim]
+ k = x[..., query_dim : 2 * query_dim]
+ # p is the position-encoding query
+ p = x[..., 2 * query_dim :]
+ assert p.shape[-1] == num_heads * pos_head_dim
+
+ # Pad cached left contexts
+ assert cached_key.shape[0] == left_context_len, (
+ cached_key.shape[0],
+ left_context_len,
+ )
+ k = torch.cat([cached_key, k], dim=0)
+ # Update cached left contexts
+ cached_key = k[-left_context_len:, ...]
+
+ # The length of key
+ k_len = k.shape[0]
+
+ q = q.reshape(seq_len, batch_size, num_heads, query_head_dim)
+ p = p.reshape(seq_len, batch_size, num_heads, pos_head_dim)
+ k = k.reshape(k_len, batch_size, num_heads, query_head_dim)
+
+ # time1 refers to target, time2 refers to source.
+ q = q.permute(2, 1, 0, 3) # (head, batch, time1, query_head_dim)
+ p = p.permute(2, 1, 0, 3) # (head, batch, time1, pos_head_dim)
+ k = k.permute(2, 1, 3, 0) # (head, batch, d_k, time2)
+
+ attn_scores = torch.matmul(q, k)
+
+ pos_emb = self.linear_pos(pos_emb)
+ seq_len2 = 2 * seq_len - 1 + left_context_len
+ pos_emb = pos_emb.reshape(-1, seq_len2, num_heads, pos_head_dim).permute(
+ 2, 0, 3, 1
+ )
+ # pos shape now: (head, {1 or batch_size}, pos_dim, seq_len2)
+
+ # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, batch, time1, seq_len2)
+ # [where seq_len2 represents relative position.]
+ pos_scores = torch.matmul(p, pos_emb)
+
+ if torch.jit.is_tracing():
+ (num_heads, batch_size, time1, n) = pos_scores.shape
+ rows = torch.arange(start=time1 - 1, end=-1, step=-1)
+ cols = torch.arange(k_len)
+ rows = rows.repeat(batch_size * num_heads).unsqueeze(-1)
+ indexes = rows + cols
+ pos_scores = pos_scores.reshape(-1, n)
+ pos_scores = torch.gather(pos_scores, dim=1, index=indexes)
+ pos_scores = pos_scores.reshape(num_heads, batch_size, time1, k_len)
+ # the following .as_strided() expression converts the last axis of pos_scores from relative
+ # to absolute position. I don't know whether I might have got the time-offsets backwards or
+ # not, but let this code define which way round it is supposed to be.
+ else:
+ pos_scores = pos_scores.as_strided(
+ (num_heads, batch_size, seq_len, k_len),
+ (
+ pos_scores.stride(0),
+ pos_scores.stride(1),
+ pos_scores.stride(2) - pos_scores.stride(3),
+ pos_scores.stride(3),
+ ),
+ storage_offset=pos_scores.stride(3) * (seq_len - 1),
+ )
+
+ attn_scores = attn_scores + pos_scores
+
+ assert attn_scores.shape == (
+ num_heads,
+ batch_size,
+ seq_len,
+ k_len,
+ ), attn_scores.shape
+
+ if key_padding_mask is not None:
+ assert key_padding_mask.shape == (batch_size, k_len), key_padding_mask.shape
+ attn_scores = attn_scores.masked_fill(
+ key_padding_mask.unsqueeze(1),
+ -1000,
+ )
+
+ attn_weights = attn_scores.softmax(dim=-1)
+
+ return attn_weights, cached_key
+
+ def _print_attn_entropy(self, attn_weights: Tensor):
+ # attn_weights: (num_heads, batch_size, seq_len, seq_len)
+ (num_heads, batch_size, seq_len, seq_len) = attn_weights.shape
+
+ with torch.no_grad():
+ with torch.cuda.amp.autocast(enabled=False):
+ attn_weights = attn_weights.to(torch.float32)
+ attn_weights_entropy = (
+ -((attn_weights + 1.0e-20).log() * attn_weights)
+ .sum(dim=-1)
+ .mean(dim=(1, 2))
+ )
+ logging.info(
+ f"name={self.name}, attn_weights_entropy = {attn_weights_entropy}"
+ )
+
+
+class SelfAttention(nn.Module):
+ """
+ The simplest possible attention module. This one works with already-computed attention
+ weights, e.g. as computed by RelPositionMultiheadAttentionWeights.
+
+ Args:
+ embed_dim: the input and output embedding dimension
+ num_heads: the number of attention heads
+ value_head_dim: the value dimension per head
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ value_head_dim: int,
+ ) -> None:
+ super().__init__()
+ self.in_proj = nn.Linear(embed_dim, num_heads * value_head_dim, bias=True)
+
+ self.out_proj = ScaledLinear(
+ num_heads * value_head_dim, embed_dim, bias=True, initial_scale=0.05
+ )
+
+ self.whiten = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(7.5, ratio=3.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ ) -> Tensor:
+ """
+ Args:
+ x: input tensor, of shape (seq_len, batch_size, embed_dim)
+ attn_weights: a tensor of shape (num_heads, batch_size, seq_len, seq_len),
+ with seq_len being interpreted as (tgt_seq_len, src_seq_len). Expect
+ attn_weights.sum(dim=-1) == 1.
+ Returns:
+ a tensor with the same shape as x.
+ """
+ (seq_len, batch_size, embed_dim) = x.shape
+ num_heads = attn_weights.shape[0]
+ assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len)
+
+ x = self.in_proj(x) # (seq_len, batch_size, num_heads * value_head_dim)
+ x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3)
+ # now x: (num_heads, batch_size, seq_len, value_head_dim)
+ value_head_dim = x.shape[-1]
+
+ # todo: see whether there is benefit in overriding matmul
+ x = torch.matmul(attn_weights, x)
+ # v: (num_heads, batch_size, seq_len, value_head_dim)
+
+ x = (
+ x.permute(2, 1, 0, 3)
+ .contiguous()
+ .view(seq_len, batch_size, num_heads * value_head_dim)
+ )
+
+ # returned value is of shape (seq_len, batch_size, embed_dim), like the input.
+ x = self.out_proj(x)
+ x = self.whiten(x)
+
+ return x
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ cached_val: Tensor,
+ left_context_len: int,
+ ) -> Tuple[Tensor, Tensor]:
+ """
+ Args:
+ x: input tensor, of shape (seq_len, batch_size, embed_dim)
+ attn_weights: a tensor of shape (num_heads, batch_size, seq_len, seq_len),
+ with seq_len being interpreted as (tgt_seq_len, src_seq_len). Expect
+ attn_weights.sum(dim=-1) == 1.
+ cached_val: cached attention value tensor of left context,
+ of shape (left_context_len, batch_size, value_dim)
+ left_context_len: number of left context frames.
+
+ Returns:
+ - attention weighted output, a tensor with the same shape as x.
+ - updated cached attention value tensor of left context.
+ """
+ (seq_len, batch_size, embed_dim) = x.shape
+ num_heads = attn_weights.shape[0]
+ seq_len2 = seq_len + left_context_len
+ assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len2)
+
+ x = self.in_proj(x) # (seq_len, batch_size, num_heads * value_head_dim)
+
+ # Pad cached left contexts
+ assert cached_val.shape[0] == left_context_len, (
+ cached_val.shape[0],
+ left_context_len,
+ )
+ x = torch.cat([cached_val, x], dim=0)
+ # Update cached left contexts
+ cached_val = x[-left_context_len:, ...]
+
+ x = x.reshape(seq_len2, batch_size, num_heads, -1).permute(2, 1, 0, 3)
+ # now x: (num_heads, batch_size, seq_len, value_head_dim)
+ value_head_dim = x.shape[-1]
+
+ # todo: see whether there is benefit in overriding matmul
+ x = torch.matmul(attn_weights, x)
+ # v: (num_heads, batch_size, seq_len, value_head_dim)
+
+ x = (
+ x.permute(2, 1, 0, 3)
+ .contiguous()
+ .view(seq_len, batch_size, num_heads * value_head_dim)
+ )
+
+ # returned value is of shape (seq_len, batch_size, embed_dim), like the input.
+ x = self.out_proj(x)
+
+ return x, cached_val
+
+
+class FeedforwardModule(nn.Module):
+ """Feedforward module in Zipformer2 model."""
+
+ def __init__(self, embed_dim: int, feedforward_dim: int, dropout: FloatLike):
+ super(FeedforwardModule, self).__init__()
+ self.in_proj = nn.Linear(embed_dim, feedforward_dim)
+
+ self.hidden_balancer = Balancer(
+ feedforward_dim,
+ channel_dim=-1,
+ min_positive=0.3,
+ max_positive=1.0,
+ min_abs=0.75,
+ max_abs=5.0,
+ )
+
+ # shared_dim=0 means we share the dropout mask along the time axis
+ self.out_proj = ActivationDropoutAndLinear(
+ feedforward_dim,
+ embed_dim,
+ activation="SwooshL",
+ dropout_p=dropout,
+ dropout_shared_dim=0,
+ bias=True,
+ initial_scale=0.1,
+ )
+
+ self.out_whiten = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(7.5),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ def forward(self, x: Tensor):
+ x = self.in_proj(x)
+ x = self.hidden_balancer(x)
+ # out_proj contains SwooshL activation, then dropout, then linear.
+ x = self.out_proj(x)
+ x = self.out_whiten(x)
+ return x
+
+
+class NonlinAttention(nn.Module):
+ """This is like the ConvolutionModule, but refactored so that we use multiplication by attention weights (borrowed
+ from the attention module) in place of actual convolution. We also took out the second nonlinearity, the
+ one after the attention mechanism.
+
+ Args:
+ channels (int): The number of channels of conv layers.
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ hidden_channels: int,
+ ) -> None:
+ super().__init__()
+
+ self.hidden_channels = hidden_channels
+
+ self.in_proj = nn.Linear(channels, hidden_channels * 3, bias=True)
+
+ # balancer that goes before the sigmoid. Have quite a large min_abs value, at 2.0,
+ # because we noticed that well-trained instances of this module have abs-value before the sigmoid
+ # starting from about 3, and poorly-trained instances of the module have smaller abs values
+ # before the sigmoid.
+ self.balancer = Balancer(
+ hidden_channels,
+ channel_dim=-1,
+ min_positive=ScheduledFloat((0.0, 0.25), (20000.0, 0.05)),
+ max_positive=ScheduledFloat((0.0, 0.75), (20000.0, 0.95)),
+ min_abs=0.5,
+ max_abs=5.0,
+ )
+ self.tanh = nn.Tanh()
+
+ self.identity1 = Identity() # for diagnostics.
+ self.identity2 = Identity() # for diagnostics.
+ self.identity3 = Identity() # for diagnostics.
+
+ self.out_proj = ScaledLinear(
+ hidden_channels, channels, bias=True, initial_scale=0.05
+ )
+
+ self.whiten1 = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(5.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ self.whiten2 = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(5.0, ratio=3.0),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ ) -> Tensor:
+ """.
+ Args:
+ x: a Tensor of shape (seq_len, batch_size, num_channels)
+ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len)
+ Returns:
+ a Tensor with the same shape as x
+ """
+ x = self.in_proj(x)
+
+ (seq_len, batch_size, _) = x.shape
+ hidden_channels = self.hidden_channels
+
+ s, x, y = x.chunk(3, dim=2)
+
+ # s will go through tanh.
+
+ s = self.balancer(s)
+ s = self.tanh(s)
+
+ s = s.unsqueeze(-1).reshape(seq_len, batch_size, hidden_channels)
+ x = self.whiten1(x)
+ x = x * s
+ x = self.identity1(x) # diagnostics only, it's the identity.
+
+ (seq_len, batch_size, embed_dim) = x.shape
+ num_heads = attn_weights.shape[0]
+ assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len)
+
+ x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3)
+ # now x: (num_heads, batch_size, seq_len, head_dim)
+ x = torch.matmul(attn_weights, x)
+ # now x: (num_heads, batch_size, seq_len, head_dim)
+ x = x.permute(2, 1, 0, 3).reshape(seq_len, batch_size, -1)
+
+ y = self.identity2(y)
+ x = x * y
+ x = self.identity3(x)
+
+ x = self.out_proj(x)
+ x = self.whiten2(x)
+ return x
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ attn_weights: Tensor,
+ cached_x: Tensor,
+ left_context_len: int,
+ ) -> Tuple[Tensor, Tensor]:
+ """.
+ Args:
+ x: a Tensor of shape (seq_len, batch_size, num_channels)
+ attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len)
+ cached_x: left context, a Tensor of shape
+ (num_heads, batch_size, left_context_len, head_dim)
+ left_context_len: number of left context frames.
+ Returns:
+ - a Tensor with the same shape as x
+ - updated left context with same shape as cached_x
+ """
+ x = self.in_proj(x)
+
+ (seq_len, batch_size, _) = x.shape
+ hidden_channels = self.hidden_channels
+
+ s, x, y = x.chunk(3, dim=2)
+
+ # s will go through tanh.
+ s = self.tanh(s)
+
+ s = s.unsqueeze(-1).reshape(seq_len, batch_size, hidden_channels)
+ x = x * s
+
+ (seq_len, batch_size, embed_dim) = x.shape
+ num_heads = attn_weights.shape[0]
+ assert attn_weights.shape == (
+ num_heads,
+ batch_size,
+ seq_len,
+ left_context_len + seq_len,
+ )
+
+ x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3)
+ # now x: (num_heads, batch_size, seq_len, head_dim)
+
+ # Pad cached tensor
+ assert cached_x.shape[2] == left_context_len, (
+ cached_x.shape[2],
+ left_context_len,
+ )
+ x_pad = torch.cat([cached_x, x], dim=2)
+ # Update cached tensor
+ cached_x = x_pad[:, :, -left_context_len:, :]
+
+ x = torch.matmul(attn_weights, x_pad)
+ # now x: (num_heads, batch_size, seq_len, head_dim)
+ x = x.permute(2, 1, 0, 3).reshape(seq_len, batch_size, -1)
+
+ x = x * y
+
+ x = self.out_proj(x)
+ return x, cached_x
+
+
+class ConvolutionModule(nn.Module):
+ """ConvolutionModule in Zipformer2 model.
+ Modified from https://github.com/espnet/espnet/blob/master/espnet/nets/pytorch_backend/zipformer/convolution.py
+
+ Args:
+ channels (int): The number of channels of conv layers.
+ kernel_size (int): Kernerl size of conv layers.
+ bias (bool): Whether to use bias in conv layers (default=True).
+
+ """
+
+ def __init__(
+ self,
+ channels: int,
+ kernel_size: int,
+ causal: bool,
+ ) -> None:
+ """Construct a ConvolutionModule object."""
+ super(ConvolutionModule, self).__init__()
+ # kernerl_size should be a odd number for 'SAME' padding
+ assert (kernel_size - 1) % 2 == 0
+
+ bottleneck_dim = channels
+ self.causal = causal
+
+ self.in_proj = nn.Linear(
+ channels,
+ 2 * bottleneck_dim,
+ )
+ # the gradients on in_proj are a little noisy, likely to do with the
+ # sigmoid in glu.
+
+ # after in_proj we put x through a gated linear unit (nn.functional.glu).
+ # For most layers the normal rms value of channels of x seems to be in the range 1 to 4,
+ # but sometimes, for some reason, for layer 0 the rms ends up being very large,
+ # between 50 and 100 for different channels. This will cause very peaky and
+ # sparse derivatives for the sigmoid gating function, which will tend to make
+ # the loss function not learn effectively. (for most layers the average absolute values
+ # are in the range 0.5..9.0, and the average p(x>0), i.e. positive proportion,
+ # at the output of pointwise_conv1.output is around 0.35 to 0.45 for different
+ # layers, which likely breaks down as 0.5 for the "linear" half and
+ # 0.2 to 0.3 for the part that goes into the sigmoid. The idea is that if we
+ # constrain the rms values to a reasonable range via a constraint of max_abs=10.0,
+ # it will be in a better position to start learning something, i.e. to latch onto
+ # the correct range.
+ self.balancer1 = Balancer(
+ bottleneck_dim,
+ channel_dim=-1,
+ min_positive=ScheduledFloat((0.0, 0.05), (8000.0, 0.025)),
+ max_positive=1.0,
+ min_abs=1.5,
+ max_abs=ScheduledFloat((0.0, 5.0), (8000.0, 10.0), default=1.0),
+ )
+
+ self.activation1 = Identity() # for diagnostics
+
+ self.sigmoid = nn.Sigmoid()
+
+ self.activation2 = Identity() # for diagnostics
+
+ assert kernel_size % 2 == 1
+
+ self.depthwise_conv = (
+ ChunkCausalDepthwiseConv1d(channels=bottleneck_dim, kernel_size=kernel_size)
+ if causal
+ else nn.Conv1d(
+ in_channels=bottleneck_dim,
+ out_channels=bottleneck_dim,
+ groups=bottleneck_dim,
+ kernel_size=kernel_size,
+ padding=kernel_size // 2,
+ )
+ )
+
+ self.balancer2 = Balancer(
+ bottleneck_dim,
+ channel_dim=1,
+ min_positive=ScheduledFloat((0.0, 0.1), (8000.0, 0.05)),
+ max_positive=1.0,
+ min_abs=ScheduledFloat((0.0, 0.2), (20000.0, 0.5)),
+ max_abs=10.0,
+ )
+
+ self.whiten = Whiten(
+ num_groups=1,
+ whitening_limit=_whitening_schedule(7.5),
+ prob=(0.025, 0.25),
+ grad_scale=0.01,
+ )
+
+ self.out_proj = ActivationDropoutAndLinear(
+ bottleneck_dim,
+ channels,
+ activation="SwooshR",
+ dropout_p=0.0,
+ initial_scale=0.05,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ src_key_padding_mask: Optional[Tensor] = None,
+ chunk_size: int = -1,
+ ) -> Tensor:
+ """Compute convolution module.
+
+ Args:
+ x: Input tensor (#time, batch, channels).
+ src_key_padding_mask: the mask for the src keys per batch (optional):
+ (batch, #time), contains True in masked positions.
+
+ Returns:
+ Tensor: Output tensor (#time, batch, channels).
+
+ """
+
+ x = self.in_proj(x) # (time, batch, 2*channels)
+
+ x, s = x.chunk(2, dim=2)
+ s = self.balancer1(s)
+ s = self.sigmoid(s)
+ x = self.activation1(x) # identity.
+ x = x * s
+ x = self.activation2(x) # identity
+
+ # (time, batch, channels)
+
+ # exchange the temporal dimension and the feature dimension
+ x = x.permute(1, 2, 0) # (#batch, channels, time).
+
+ if src_key_padding_mask is not None:
+ x = x.masked_fill(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
+
+ if (
+ not torch.jit.is_scripting()
+ and not torch.jit.is_tracing()
+ and chunk_size >= 0
+ ):
+ # Not support exporting a model for simulated streaming decoding
+ assert (
+ self.causal
+ ), "Must initialize model with causal=True if you use chunk_size"
+ x = self.depthwise_conv(x, chunk_size=chunk_size)
+ else:
+ x = self.depthwise_conv(x)
+
+ x = self.balancer2(x)
+ x = x.permute(2, 0, 1) # (time, batch, channels)
+
+ x = self.whiten(x) # (time, batch, channels)
+ x = self.out_proj(x) # (time, batch, channels)
+
+ return x
+
+ def streaming_forward(
+ self,
+ x: Tensor,
+ cache: Tensor,
+ src_key_padding_mask: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ """Compute convolution module in streaming forward mode.
+
+ Args:
+ x: Input tensor (#time, batch, channels).
+ cache: cached left context for depthwise_conv of shape
+ (#batch, channels, left_pad)
+ src_key_padding_mask: the mask for the src keys per batch (optional):
+ (batch, #time), contains True in masked positions.
+
+ Returns:
+ - Output tensor (#time, batch, channels).
+ - Updated cache (#batch, channels, left_pad)
+ """
+
+ x = self.in_proj(x) # (time, batch, 2*channels)
+
+ x, s = x.chunk(2, dim=2)
+ s = self.sigmoid(s)
+ x = x * s
+ # (time, batch, channels)
+
+ # exchange the temporal dimension and the feature dimension
+ x = x.permute(1, 2, 0) # (#batch, channels, time).
+
+ if src_key_padding_mask is not None:
+ x = x.masked_fill(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0)
+
+ x, cache = self.depthwise_conv.streaming_forward(x, cache=cache)
+
+ x = x.permute(2, 0, 1) # (time, batch, channels)
+
+ x = self.out_proj(x) # (time, batch, channels)
+
+ return x, cache
+
+
+class ScalarMultiply(nn.Module):
+ def __init__(self, scale: float):
+ super().__init__()
+ self.scale = scale
+
+ def forward(self, x):
+ return x * self.scale
+
+
+def _test_zipformer_main(causal: bool = False):
+ batch_size = 5
+ seq_len = 20
+ # Just make sure the forward pass runs.
+
+ c = Zipformer2(
+ encoder_dim=(64, 96),
+ encoder_unmasked_dim=(48, 64),
+ num_heads=(4, 4),
+ causal=causal,
+ chunk_size=(4,) if causal else (-1,),
+ left_context_frames=(64,),
+ )
+ batch_size = 5
+ seq_len = 20
+ # Just make sure the forward pass runs.
+ f = c(
+ torch.randn(seq_len, batch_size, 64),
+ torch.full((batch_size,), seq_len, dtype=torch.int64),
+ )
+ f[0].sum().backward()
+ c.eval()
+ f = c(
+ torch.randn(seq_len, batch_size, 64),
+ torch.full((batch_size,), seq_len, dtype=torch.int64),
+ )
+ f # to remove flake8 warnings
+
+
+if __name__ == "__main__":
+ logging.getLogger().setLevel(logging.INFO)
+ torch.set_num_threads(1)
+ torch.set_num_interop_threads(1)
+ _test_zipformer_main(False)
+ _test_zipformer_main(True)
diff --git a/egs/europarl_st/SRT/local/README.md b/egs/europarl_st/SRT/local/README.md
new file mode 100644
index 0000000000..68ab9137b7
--- /dev/null
+++ b/egs/europarl_st/SRT/local/README.md
@@ -0,0 +1,422 @@
+# Europarl-ST Preprocessing Scripts
+
+A complete preprocessing pipeline for the [Europarl-ST](https://www.mllp.upv.es/europarl-st/) dataset, converting raw speech translation data into training-ready Lhotse CutSet manifests with pre-extracted FBANK features.
+
+## Overview
+
+The pipeline transforms raw Europarl-ST data through the following stages:
+
+```
+Raw Europarl-ST v1.1
+ │
+ ▼
+┌─────────────────┐
+│ org_to_jsonl.py │ Extract audio segments + build per-language-pair JSONL
+└────────┬────────┘
+ │
+ ▼
+┌──────────────────────┐
+│ normalize_texts.py │ Apply Whisper-style text normalization
+└────────┬─────────────┘
+ │
+ ▼
+┌──────────────────┐
+│ texts_to_cuts.py │ Generate Lhotse CutSet manifests + extract FBANK features
+└────────┬─────────┘
+ │
+ ▼
+┌────────────────────────┐
+│ filter_cuts_texts.py │ Remove entries with empty text/st_text
+└────────┬───────────────┘
+ │
+ ▼
+┌──────────────────────┐
+│ check_manifests.py │ Validate final manifests
+└──────────────────────┘
+
+┌──────────────────┐
+│ train_bpe.py │ Train SentencePiece BPE model + generate tokens.txt
+└──────────────────┘
+```
+
+## Prerequisites
+
+### Optional
+
+- [openai-whisper](https://github.com/openai/whisper) (`pip install openai-whisper`) — provides official text normalizers. If not installed, the scripts fall back to built-in local normalizers.
+
+## Directory Structure
+
+After running the full pipeline, the expected layout is:
+
+```
+Europarl-ST/
+├── v1.1/ # Raw Europarl-ST data (input)
+│ ├── es/
+│ ├── de/
+│ ├── en/
+│ └── ...
+├── audio/ # Extracted FLAC audio segments
+│ ├── train/
+│ ├── valid/
+│ └── test/
+├── texts/ # Per-language-pair JSONL files
+│ ├── es_de/
+│ ├── es_en/
+│ └── ...
+├── normalizer/ # Normalized JSONL files
+│ ├── es_de/
+│ ├── es_en/
+│ └── ...
+├── fbank/ # FBANK feature storage (.lca chunks)
+│ └── feature_cache.json
+├── manifests/ # Final Lhotse CutSet manifests
+│ ├── es_de/
+│ ├── es_en/
+│ └── ...
+├── bpe/ # Trained BPE models and token lists
+│ ├── asr9/ # ASR BPE (9-language shared)
+│ │ ├── bpe.model
+│ │ └── tokens.txt
+│ ├── ast9/ # ST BPE (9-language shared)
+│ │ ├── bpe.model
+│ │ └── tokens.txt
+│ └── ...
+└── local/ # This directory
+ ├── README.md
+ ├── org_to_jsonl.py
+ ├── normalize_texts.py
+ ├── normalize_jsonl_with_whisper.py
+ ├── filter_cuts_texts.py
+ ├── texts_to_cuts.py
+ ├── check_manifests.py
+ └── train_bpe.py
+```
+
+## Scripts
+
+### 1. `org_to_jsonl.py`
+
+Extracts audio segments from raw Europarl-ST and produces per-language-pair JSONL files.
+
+**What it does:**
+- Iterates over all 9 languages (es, de, en, fr, nl, pl, pt, ro, it) and their pair combinations
+- Cuts audio segments from `.m4a` files based on timestamps, converts to `.flac`
+- Remaps dataset splits: `train` → `train`, `dev` → `valid`, `test` → `test`
+- Writes JSONL entries with fields: `source`, `duration`, `text`, `st_text`
+
+**Usage:**
+```bash
+python org_to_jsonl.py \
+ --data-dir ../v1.1 \
+ --output-dir ../audio
+```
+
+**Arguments:**
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--data-dir` | `../v1.1` | Path to the raw Europarl-ST v1.1 directory |
+| `--output-dir` | `../audio` | Directory for converted FLAC segments |
+
+---
+
+### 2. `normalize_texts.py`
+
+Applies Whisper-style text normalization to the JSONL files produced by `org_to_jsonl.py`.
+
+**What it does:**
+- Normalizes Unicode (NFKC), replaces fancy quotes/dashes, removes control characters
+- Automatically detects English fields from directory names (e.g., `es_en`) and applies the stricter English normalizer (retains only `[a-z0-9']`)
+- Uses official `whisper.normalizers` if available, otherwise falls back to built-in implementations
+
+**Usage:**
+```bash
+python normalize_texts.py \
+ --src-dir ../texts \
+ --dst-dir ../normalizer \
+ --fields text st_text \
+ --normalizer basic \
+ --skip-existing
+```
+
+**Arguments:**
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--src-dir` | `../texts` | Root directory containing JSONL files |
+| `--dst-dir` | `../texts/normalizer` | Destination for normalized output |
+| `--fields` | `text st_text` | JSON keys to normalize (supports dot paths) |
+| `--normalizer` | `basic` | Base normalizer: `basic` or `english` |
+| `--no-auto-english` | (disabled) | Disable auto-detection of English fields |
+| `--skip-existing` | (disabled) | Skip already-normalized files |
+| `--dry-run` | (disabled) | Report stats without writing |
+| `--verbose` | (disabled) | Enable DEBUG logging |
+
+---
+
+### 3. `normalize_jsonl_with_whisper.py`
+
+Similar to `normalize_texts.py`, but operates on Lhotse CutSet-format `.jsonl.gz` files (post `texts_to_cuts.py`).
+
+**What it does:**
+- Normalizes `supervision.text` based on the `language` field
+- Normalizes `supervision.custom.st_text` based on the `custom.lang` field
+- Selects English normalizer for English fields, basic normalizer for others
+
+**Usage:**
+```bash
+python normalize_jsonl_with_whisper.py \
+ --input /path/to/input.jsonl.gz \
+ --output /path/to/output.jsonl.gz
+```
+
+**Arguments:**
+| Argument | Description |
+|----------|-------------|
+| `--input` | Path to the source `.jsonl.gz` file |
+| `--output` | Path to the destination `.jsonl.gz` file |
+| `--keep-empty` | Keep empty lines from input (skipped by default) |
+
+---
+
+### 4. `texts_to_cuts.py`
+
+Converts normalized JSONL files into Lhotse CutSet manifests with pre-extracted FBANK features.
+
+**What it does:**
+- Reads normalized JSONL, resolves audio paths, builds `Recording` and `SupervisionSegment` objects
+- Extracts 80-dim FBANK features using `kaldifeat` (GPU-accelerated when available)
+- Stores features in LilcomChunky (`.lca`) format for efficient random access
+- Maintains a feature cache to avoid redundant computation across runs
+- Supports sharding for large training sets
+
+**Usage:**
+```bash
+python texts_to_cuts.py \
+ --src-dir ../normalizer \
+ --dst-dir ../manifests \
+ --audio-root /path/to/audio/root \
+ --storage-root ../fbank \
+ --feature-cache ../fbank/feature_cache.json \
+ --num-workers 8 \
+ --batch-duration 600 \
+ --skip-missing-audio \
+ --verbose
+```
+
+**Arguments:**
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--src-dir` | `../texts/normalizer` | Normalized JSONL root directory |
+| `--dst-dir` | `../cut_manifests` | Output directory for CutSet manifests |
+| `--audio-root` | (parent of dataset dir) | Base path for resolving relative audio paths |
+| `--storage-root` | `../fbank_storage` | Directory for `.lca` feature chunks |
+| `--num-workers` | `8` | Workers for parallel feature extraction |
+| `--batch-duration` | `600.0` | Total audio seconds per extraction batch |
+| `--device` | `auto` | Device: `auto`, `cpu`, or `cuda` |
+| `--train-shard-duration` | `0` (disabled) | Max audio seconds per train shard |
+| `--skip-missing-audio` | (disabled) | Skip missing audio instead of raising |
+| `--feature-cache` | (in storage root) | Path to feature cache JSON |
+| `--refresh-cache` | (disabled) | Ignore cache, recompute all features |
+| `--overwrite` | (disabled) | Overwrite existing outputs |
+| `--verbose` | (disabled) | Enable DEBUG logging |
+
+---
+
+### 5. `filter_cuts_texts.py`
+
+Removes CutSet entries that have empty `text` or `st_text` fields.
+
+**What it does:**
+- Scans all `*_cuts.jsonl` and `*_cuts.jsonl.gz` files in the manifest directory
+- Removes any cut whose supervision has an empty `text` or `custom.st_text`
+- Supports in-place overwrite or output to a separate directory
+
+**Usage:**
+```bash
+python filter_cuts_texts.py \
+ --manifest-dir ../manifests \
+ --output-dir ../manifests \
+ --overwrite \
+ --verbose
+```
+
+**Arguments:**
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--manifest-dir` | `../manifests` | Directory containing CutSet manifests |
+| `--output-dir` | (same as manifest-dir) | Output directory for filtered manifests |
+| `--overwrite` | (disabled) | Allow in-place replacement |
+| `--dry-run` | (disabled) | Report stats without writing |
+| `--verbose` | (disabled) | Enable DEBUG logging |
+
+---
+
+### 6. `check_manifests.py`
+
+Validates the final CutSet manifests for correctness and data integrity.
+
+**What it does:**
+- Runs Lhotse's built-in `validate()` on all `*.jsonl.gz` manifests
+- Optionally loads actual audio/features to detect stale metadata (`--read-data`)
+- Reports per-cut failure counts and sample IDs (`--per-cut-details`)
+- Supports multi-process parallel validation
+
+**Usage:**
+```bash
+python check_manifests.py \
+ --manifests-dir ../manifests \
+ --num-workers 8 \
+ --read-data \
+ --verbose
+```
+
+**Arguments:**
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--manifests-dir` | `../manifests` | Directory containing manifests |
+| `--pattern` | `*.jsonl.gz` | Glob pattern for files to validate |
+| `--limit` | (none) | Cap number of files to inspect |
+| `--num-workers` | `8` | Parallel worker processes |
+| `--read-data` | (disabled) | Load audio/features (slow but thorough) |
+| `--per-cut-details` | (disabled) | Report per-cut failure counts |
+| `--bad-cut-samples` | `5` | Number of failing cut IDs to show |
+| `--verbose` | (disabled) | Enable verbose logging |
+
+### 7. `train_bpe.py`
+
+Trains a SentencePiece BPE model from text transcripts and generates a `tokens.txt` vocabulary file.
+
+**What it does:**
+- Trains a unigram SentencePiece model with user-defined special symbols (blank, sos/eos, language tags)
+- Generates `bpe.model` and `tokens.txt` in the specified output directory
+- Supports custom vocabulary size for different tokenization granularities
+
+**Usage:**
+```bash
+python train_bpe.py \
+ --lang-dir ../Europarl-ST/bpe/ast9 \
+ --transcript /path/to/training_text.txt \
+ --vocab-size 6000
+```
+
+**Arguments:**
+| Argument | Description |
+|----------|-------------|
+| `--lang-dir` | Output directory for `bpe.model` and `tokens.txt` |
+| `--transcript` | Path to training text file (one sentence per line) |
+| `--vocab-size` | Target vocabulary size |
+
+**Notes:**
+- The model includes pre-defined language tags: `<2en>`, `<2de>`, `<2es>`, `<2fr>`, `<2it>`, `<2nl>`, `<2pl>`, `<2pt>`, `<2ro>`
+- Special tokens `` and `` are reserved
+- Requires `sentencepiece >= 0.1.96` (`pip install sentencepiece`)
+
+---
+
+## Supported Languages
+
+The dataset covers 9 European languages:
+
+| Code | Language |
+|------|----------|
+| `es` | Spanish |
+| `de` | German |
+| `en` | English |
+| `fr` | French |
+| `nl` | Dutch |
+| `pl` | Polish |
+| `pt` | Portuguese |
+| `ro` | Romanian |
+| `it` | Italian |
+
+All non-identical language pairs are processed (72 pairs total).
+
+## Output Format
+
+Each final CutSet manifest entry (MonoCut) contains:
+
+```json
+{
+ "id": "en_5525-758",
+ "start": 0,
+ "duration": 6.76,
+ "channel": 0,
+ "supervisions": [
+ {
+ "id": "en_5525",
+ "recording_id": "en_5525",
+ "start": 0.0,
+ "duration": 6.76,
+ "channel": 0,
+ "text": "mister president one of the key issues in the new lisbon treaty is the increased role of the european union in the world",
+ "language": "en",
+ "speaker": "unknown",
+ "custom": {
+ "st_text": "herr präsident einer der hauptpunkte des neuen vertrages von lissabon ist die aufwertung der rolle der europäischen union in der welt",
+ "lang": "de"
+ }
+ }
+ ],
+ "recording": {
+ "id": "en_5525",
+ "sources": [
+ {
+ "type": "file",
+ "channels": [0],
+ "source": "data/Europarl-ST/audio/test/en_5525.flac"
+ }
+ ],
+ "sampling_rate": 16000,
+ "num_samples": 108160,
+ "duration": 6.76
+ },
+ "type": "MonoCut"
+}
+```
+
+Key fields:
+- `supervisions[0].text`: ASR ground truth transcription
+- `supervisions[0].language`: source language
+- `supervisions[0].custom.st_text`: ST ground truth translation
+- `supervisions[0].custom.lang`: target language
+
+## Quick Start
+
+Run the full pipeline from the dataset root:
+
+```bash
+cd /path/to/Europarl-ST
+
+# Step 1: Extract audio and build JSONL
+python local/org_to_jsonl.py --data-dir ./v1.1 --output-dir ./audio
+
+# Step 2: Normalize text
+python local/normalize_texts.py --src-dir ./texts --dst-dir ./normalizer
+
+# Step 3: Generate CutSet manifests with FBANK features
+python local/texts_to_cuts.py \
+ --src-dir ./normalizer \
+ --dst-dir ./manifests \
+ --audio-root /path/to/audio/root \
+ --storage-root ./fbank \
+ --num-workers 8 \
+ --skip-missing-audio
+
+# Step 4: Filter out entries with empty text
+python local/filter_cuts_texts.py \
+ --manifest-dir ./manifests \
+ --overwrite
+
+# Step 5: Validate
+python local/check_manifests.py --manifests-dir ./manifests --read-data
+
+# Step 6: Train BPE model
+python local/train_bpe.py \
+ --lang-dir ./bpe/ast9 \
+ --transcript /path/to/all_training_text.txt \
+ --vocab-size 6000
+```
+
+## License
+
+Apache License 2.0
diff --git a/egs/europarl_st/SRT/local/check_manifests.py b/egs/europarl_st/SRT/local/check_manifests.py
new file mode 100644
index 0000000000..91505de27d
--- /dev/null
+++ b/egs/europarl_st/SRT/local/check_manifests.py
@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Utility script to sanity-check Europarl-ST CutSet manifests.
+
+It walks through all *.jsonl.gz files under the provided directory,
+runs Lhotse's built-in validators, and reports which manifests fail.
+
+Optionally it can read the underlying audio/features (`--read-data`)
+to catch mismatches between metadata and stored tensors (slow but thorough),
+and with `--per-cut-details` it will report how many individual cuts fail per manifest.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+from lhotse import CutSet
+from lhotse.qa import validate
+
+DEFAULT_MANIFEST_ROOT = Path(__file__).resolve().parent.parent / "manifests"
+
+
+@dataclass
+class ValidationResult:
+ path: Path
+ status: str
+ num_cuts: Optional[int] = None
+ error: Optional[str] = None
+ bad_cut_count: Optional[int] = None
+ bad_cut_samples: Optional[List[str]] = None
+
+ @property
+ def ok(self) -> bool:
+ return self.status == "ok"
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Validate Europarl-ST CutSet manifests produced by texts_to_cuts.py"
+ )
+ parser.add_argument(
+ "--manifests-dir",
+ type=Path,
+ default=DEFAULT_MANIFEST_ROOT,
+ help="Directory that contains *.jsonl.gz CutSet manifests.",
+ )
+ parser.add_argument(
+ "--pattern",
+ type=str,
+ default="*.jsonl.gz",
+ help="Glob-style pattern (relative to --manifests-dir) for files to validate.",
+ )
+ parser.add_argument(
+ "--limit",
+ type=int,
+ default=None,
+ help="Optionally cap the number of files to inspect (useful for quick smoke tests).",
+ )
+ parser.add_argument(
+ "--num-workers",
+ type=int,
+ default=8,
+ help="How many parallel worker processes to spawn. Reading features/audio is I/O bound, "
+ "so >1 workers can shorten runtime.",
+ )
+ parser.add_argument(
+ "--read-data",
+ action="store_true",
+ help=(
+ "Load audio/features referenced by the manifests while validating. "
+ "This is significantly slower but detects stale feature caches."
+ ),
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Enable verbose logging (per-file progress).",
+ )
+ parser.add_argument(
+ "--per-cut-details",
+ action="store_true",
+ help=(
+ "Iterate through every cut and count the number of failures per manifest. "
+ "When combined with --read-data this reveals exactly how many cuts have stale features."
+ ),
+ )
+ parser.add_argument(
+ "--bad-cut-samples",
+ type=int,
+ default=5,
+ help="When --per-cut-details is enabled, how many failing cut IDs to show per manifest.",
+ )
+ return parser.parse_args()
+
+
+def find_manifest_files(root: Path, pattern: str) -> List[Path]:
+ files = sorted(root.rglob(pattern))
+ return files
+
+
+def _collect_bad_cuts(
+ cuts: CutSet, read_data: bool, sample_limit: int
+) -> Tuple[int, List[str], Optional[str]]:
+ bad_count = 0
+ samples: List[str] = []
+ first_error: Optional[str] = None
+ for cut in cuts:
+ try:
+ validate(cut, read_data=read_data)
+ except AssertionError as exc:
+ bad_count += 1
+ msg = str(exc)
+ if first_error is None:
+ first_error = msg
+ if len(samples) < sample_limit:
+ samples.append(f"{cut.id}: {msg}")
+ return bad_count, samples, first_error
+
+
+def validate_manifest(
+ path: Path, read_data: bool, per_cut_details: bool, sample_limit: int
+) -> ValidationResult:
+ cuts: Optional[CutSet] = None
+ try:
+ cuts = CutSet.from_file(path)
+ num_cuts = len(cuts)
+ if per_cut_details:
+ bad_count, samples, first_error = _collect_bad_cuts(
+ cuts=cuts, read_data=read_data, sample_limit=sample_limit
+ )
+ if bad_count > 0:
+ return ValidationResult(
+ path=path,
+ status="invalid",
+ num_cuts=num_cuts,
+ error=first_error or "Per-cut validation failed.",
+ bad_cut_count=bad_count,
+ bad_cut_samples=samples,
+ )
+ validate(cuts, read_data=read_data)
+ return ValidationResult(path=path, status="ok", num_cuts=num_cuts)
+ except AssertionError as exc:
+ return ValidationResult(
+ path=path,
+ status="invalid",
+ num_cuts=len(cuts) if cuts is not None else None,
+ error=str(exc),
+ )
+ except Exception as exc: # pylint: disable=broad-except
+ return ValidationResult(
+ path=path,
+ status="error",
+ num_cuts=len(cuts) if cuts is not None else None,
+ error=f"{type(exc).__name__}: {exc}",
+ )
+
+
+def _validate_worker_wrapper(
+ path_str: str, read_data: bool, per_cut_details: bool, sample_limit: int
+) -> ValidationResult:
+ path = Path(path_str)
+ return validate_manifest(
+ path,
+ read_data=read_data,
+ per_cut_details=per_cut_details,
+ sample_limit=sample_limit,
+ )
+
+
+def run_serial(
+ files: List[Path],
+ read_data: bool,
+ verbose: bool,
+ per_cut_details: bool,
+ sample_limit: int,
+) -> List[ValidationResult]:
+ results: List[ValidationResult] = []
+ for idx, path in enumerate(files, start=1):
+ if verbose:
+ logging.info("Validating [%d/%d]: %s", idx, len(files), path)
+ results.append(
+ validate_manifest(
+ path=path,
+ read_data=read_data,
+ per_cut_details=per_cut_details,
+ sample_limit=sample_limit,
+ )
+ )
+ return results
+
+
+def run_parallel(
+ files: List[Path],
+ read_data: bool,
+ verbose: bool,
+ num_workers: int,
+ per_cut_details: bool,
+ sample_limit: int,
+) -> List[ValidationResult]:
+ results: List[ValidationResult] = []
+ order = {path: idx for idx, path in enumerate(files)}
+ with ProcessPoolExecutor(max_workers=num_workers) as pool:
+ futures = {
+ pool.submit(
+ _validate_worker_wrapper,
+ str(path),
+ read_data,
+ per_cut_details,
+ sample_limit,
+ ): path
+ for path in files
+ }
+ for idx, future in enumerate(as_completed(futures), start=1):
+ path = futures[future]
+ try:
+ result = future.result()
+ except Exception as exc: # pylint: disable=broad-except
+ result = ValidationResult(
+ path=path, status="error", error=f"Worker crashed: {exc}"
+ )
+ results.append(result)
+ if verbose:
+ logging.info(
+ "Validated [%d/%d]: %s (%s)", idx, len(files), path, result.status
+ )
+ # Preserve original ordering for readability
+ results.sort(key=lambda res: order.get(res.path, len(order)))
+ return results
+
+
+def summarize(results: List[ValidationResult]) -> None:
+ total = len(results)
+ ok = sum(res.ok for res in results)
+ invalid = [res for res in results if res.status == "invalid"]
+ errored = [res for res in results if res.status == "error"]
+
+ logging.info("Checked %d manifest files.", total)
+ logging.info(" OK : %d", ok)
+ logging.info(" Invalid : %d", len(invalid))
+ logging.info(" Errors : %d", len(errored))
+
+ if invalid:
+ logging.warning("First %d invalid files:", min(len(invalid), 20))
+ for res in invalid[:20]:
+ extra = ""
+ if res.bad_cut_count is not None:
+ extra = f" | bad cuts: {res.bad_cut_count}"
+ if res.bad_cut_samples:
+ sample_preview = "; ".join(res.bad_cut_samples)
+ extra += f" | samples: {sample_preview}"
+ logging.warning(" %s -> %s%s", res.path, res.error, extra)
+ if errored:
+ logging.error("First %d error files:", min(len(errored), 20))
+ for res in errored[:20]:
+ logging.error(" %s -> %s", res.path, res.error)
+
+
+def main() -> None:
+ args = parse_args()
+ logging.basicConfig(
+ format="%(asctime)s %(levelname)s %(message)s",
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ )
+
+ manifests_dir = args.manifests_dir.resolve()
+ if not manifests_dir.is_dir():
+ raise FileNotFoundError(f"Manifests directory not found: {manifests_dir}")
+
+ files = find_manifest_files(manifests_dir, args.pattern)
+ if args.limit is not None:
+ files = files[: args.limit]
+
+ if not files:
+ logging.warning(
+ "No manifests matching pattern '%s' under %s", args.pattern, manifests_dir
+ )
+ return
+
+ logging.info("Found %d manifest files to validate.", len(files))
+ if args.num_workers > 1:
+ results = run_parallel(
+ files=files,
+ read_data=args.read_data,
+ verbose=args.verbose,
+ num_workers=args.num_workers,
+ per_cut_details=args.per_cut_details,
+ sample_limit=args.bad_cut_samples,
+ )
+ else:
+ results = run_serial(
+ files=files,
+ read_data=args.read_data,
+ verbose=args.verbose,
+ per_cut_details=args.per_cut_details,
+ sample_limit=args.bad_cut_samples,
+ )
+
+ summarize(results)
+
+
+if __name__ == "__main__":
+ main()
+
+"""
+Example usage:
+
+python check_manifests.py \
+ --manifests-dir ./manifests \
+ --num-workers 8 \
+ --read-data \
+ --verbose
+"""
diff --git a/egs/europarl_st/SRT/local/filter_cuts_texts.py b/egs/europarl_st/SRT/local/filter_cuts_texts.py
new file mode 100644
index 0000000000..e4fa675f96
--- /dev/null
+++ b/egs/europarl_st/SRT/local/filter_cuts_texts.py
@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Filter CutSet manifests by removing entries whose text or st_text is empty.
+
+Usage example:
+
+python filter_cuts_texts.py \
+ --manifest-dir ./manifests \
+ --output-dir ./manifests \
+ --overwrite \
+ --verbose
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+import logging
+import shutil
+from pathlib import Path
+from typing import Iterable
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description=(
+ "Remove cuts from Lhotse CutSet manifests when any supervision has "
+ "an empty text or st_text field."
+ )
+ )
+ parser.add_argument(
+ "--manifest-dir",
+ type=Path,
+ default=Path(__file__).resolve().parent.parent / "manifests",
+ help="Directory containing *_cuts.jsonl or *_cuts.jsonl.gz files.",
+ )
+ parser.add_argument(
+ "--output-dir",
+ type=Path,
+ default=None,
+ help=(
+ "Where to write filtered manifests. Defaults to manifest-dir "
+ "(overwriting requires --overwrite)."
+ ),
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Allow rewriting files in-place (output-dir == manifest-dir).",
+ )
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ help="Do not write files, only report statistics.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Enable debug-level logging.",
+ )
+ return parser.parse_args()
+
+
+def iter_manifest_files(manifest_dir: Path) -> Iterable[Path]:
+ for path in sorted(manifest_dir.rglob("*.jsonl*")):
+ if path.is_file() and path.name.endswith("_cuts.jsonl"):
+ yield path
+ elif path.is_file() and path.name.endswith("_cuts.jsonl.gz"):
+ yield path
+
+
+def load_lines(path: Path):
+ opener = gzip.open if path.suffix == ".gz" else open
+ mode = "rt"
+ with opener(path, mode, encoding="utf-8") as f: # type: ignore
+ for line in f:
+ line = line.strip()
+ if line:
+ yield line
+
+
+def dump_lines(path: Path, lines: Iterable[str]) -> None:
+ opener = gzip.open if path.suffix == ".gz" else open
+ mode = "wt"
+ with opener(path, mode, encoding="utf-8") as f: # type: ignore
+ for line in lines:
+ f.write(line + "\n")
+
+
+def has_empty_text(cut_obj: dict) -> bool:
+ supervisions = cut_obj.get("supervisions") or []
+ if not supervisions:
+ return True
+ for supervision in supervisions:
+ text = (supervision.get("text") or "").strip()
+ st_text = (supervision.get("custom", {}).get("st_text") or "").strip()
+ if not text or not st_text:
+ return True
+ return False
+
+
+def process_manifest(
+ src_path: Path,
+ dst_path: Path,
+ dry_run: bool = False,
+) -> tuple[int, int]:
+ total = 0
+ kept = 0
+ filtered_lines = []
+
+ for line in load_lines(src_path):
+ total += 1
+ obj = json.loads(line)
+ if has_empty_text(obj):
+ continue
+ kept += 1
+ if not dry_run:
+ filtered_lines.append(json.dumps(obj, ensure_ascii=False))
+
+ if not dry_run:
+ dst_path.parent.mkdir(parents=True, exist_ok=True)
+ dump_lines(dst_path, filtered_lines)
+
+ return total, kept
+
+
+def maybe_copy(src: Path, dst: Path) -> None:
+ if src == dst:
+ return
+ dst.parent.mkdir(parents=True, exist_ok=True)
+ shutil.copy2(src, dst)
+
+
+def main() -> None:
+ args = parse_args()
+ logging.basicConfig(
+ format="%(asctime)s %(levelname)s %(message)s",
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ )
+ manifest_dir = args.manifest_dir.resolve()
+ output_dir = (args.output_dir or manifest_dir).resolve()
+
+ if output_dir == manifest_dir and not args.overwrite and not args.dry_run:
+ raise ValueError(
+ "Output directory matches manifest directory. "
+ "Use --overwrite to allow in-place replacement, "
+ "or specify --output-dir elsewhere."
+ )
+
+ logging.info("Scanning manifests under %s", manifest_dir)
+ processed = 0
+ removed = 0
+
+ for src_path in iter_manifest_files(manifest_dir):
+ rel = src_path.relative_to(manifest_dir)
+ dst_path = output_dir / rel
+ logging.info("Filtering %s -> %s", src_path, dst_path)
+ total, kept = process_manifest(
+ src_path=src_path,
+ dst_path=dst_path,
+ dry_run=args.dry_run,
+ )
+ processed += 1
+ removed += total - kept
+ logging.info(
+ "Finished %s (total=%d, kept=%d, removed=%d)",
+ src_path.name,
+ total,
+ kept,
+ total - kept,
+ )
+
+ logging.info(
+ "Done. files=%d, removed=%d entries. Dry-run=%s",
+ processed,
+ removed,
+ args.dry_run,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/europarl_st/SRT/local/normalize_jsonl_with_whisper.py b/egs/europarl_st/SRT/local/normalize_jsonl_with_whisper.py
new file mode 100644
index 0000000000..8c0018dbb7
--- /dev/null
+++ b/egs/europarl_st/SRT/local/normalize_jsonl_with_whisper.py
@@ -0,0 +1,193 @@
+#!/usr/bin/env python3
+
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Normalize ASR/AST JSONL.gz datasets with Whisper's text normalization.
+
+For every supervision entry, the script normalizes:
+ * `text` using the language specified by `language`
+ * `custom.st_text` using the language specified by `custom.lang`
+
+Usage example:
+
+python normalize_jsonl_with_whisper.py \
+ --input /path/to/input.jsonl.gz \
+ --output /path/to/output.jsonl.gz
+"""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+import sys
+from pathlib import Path
+from typing import Any, Dict, Tuple
+
+from whisper.normalizers.basic import BasicTextNormalizer
+from whisper.normalizers.english import EnglishTextNormalizer
+
+ENGLISH_ALIASES = {"en", "eng", "english"}
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description="Normalize text and st_text fields in JSONL.gz using Whisper normalizers."
+ )
+ parser.add_argument(
+ "--input", required=True, help="Path to the source .jsonl.gz file."
+ )
+ parser.add_argument(
+ "--output",
+ required=True,
+ help="Path to the destination .jsonl.gz file with normalized text.",
+ )
+ parser.add_argument(
+ "--keep-empty",
+ action="store_true",
+ help="Keep empty lines from the input (they are skipped by default).",
+ )
+ return parser.parse_args()
+
+
+def canonicalize_lang(lang: str | None) -> str | None:
+ if not lang:
+ return None
+ stripped = lang.strip().lower().replace("_", "-")
+ if "-" in stripped:
+ stripped = stripped.split("-", 1)[0]
+ return stripped
+
+
+def choose_normalizer(
+ lang: str | None,
+ english_normalizer: EnglishTextNormalizer,
+ default_normalizer: BasicTextNormalizer,
+):
+ normalized_lang = canonicalize_lang(lang)
+ if normalized_lang in ENGLISH_ALIASES:
+ return english_normalizer
+ return default_normalizer
+
+
+def normalize_text(
+ text: Any,
+ lang: str | None,
+ english_normalizer: EnglishTextNormalizer,
+ default_normalizer: BasicTextNormalizer,
+) -> Tuple[Any, bool]:
+ if not isinstance(text, str):
+ return text, False
+
+ text_stripped = text.strip()
+ if not text_stripped:
+ return text, False
+
+ normalizer = choose_normalizer(lang, english_normalizer, default_normalizer)
+ normalized_text = normalizer(text_stripped).strip()
+
+ # Whisper normalizers collapse whitespace; ensure we don't reintroduce leading/trailing spaces.
+ if normalized_text != text:
+ return normalized_text, True
+ return text, False
+
+
+def normalize_record(
+ record: Dict[str, Any],
+ english_normalizer: EnglishTextNormalizer,
+ default_normalizer: BasicTextNormalizer,
+) -> Tuple[Dict[str, Any], Dict[str, int]]:
+ stats = {"text": 0, "st_text": 0}
+ supervisions = record.get("supervisions") or []
+
+ for supervision in supervisions:
+ lang = supervision.get("language")
+ text = supervision.get("text")
+ normalized, changed = normalize_text(
+ text, lang, english_normalizer, default_normalizer
+ )
+ if changed:
+ supervision["text"] = normalized
+ stats["text"] += 1
+
+ custom = supervision.get("custom")
+ if isinstance(custom, dict) and "st_text" in custom:
+ st_lang = custom.get("lang")
+ st_text = custom.get("st_text")
+ normalized_st, st_changed = normalize_text(
+ st_text, st_lang, english_normalizer, default_normalizer
+ )
+ if st_changed:
+ custom["st_text"] = normalized_st
+ stats["st_text"] += 1
+
+ return record, stats
+
+
+def main() -> None:
+ args = parse_args()
+ input_path = Path(args.input)
+ output_path = Path(args.output)
+
+ if not input_path.is_file():
+ print(f"[ERROR] Input file not found: {input_path}", file=sys.stderr)
+ sys.exit(1)
+
+ english_normalizer = EnglishTextNormalizer()
+ default_normalizer = BasicTextNormalizer()
+
+ total_lines = 0
+ text_updates = 0
+ st_text_updates = 0
+
+ with gzip.open(input_path, "rt", encoding="utf-8") as reader, gzip.open(
+ output_path, "wt", encoding="utf-8"
+ ) as writer:
+ for line in reader:
+ total_lines += 1
+
+ stripped = line.strip("\n")
+ if not stripped:
+ if args.keep_empty:
+ writer.write("\n")
+ continue
+
+ try:
+ record = json.loads(stripped)
+ except json.JSONDecodeError as exc:
+ raise ValueError(f"Invalid JSON on line {total_lines}: {exc}") from exc
+
+ record, stats = normalize_record(
+ record, english_normalizer, default_normalizer
+ )
+ text_updates += stats["text"]
+ st_text_updates += stats["st_text"]
+
+ writer.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+ print(
+ f"Processed {total_lines} lines. "
+ f"Normalized text fields: {text_updates}, st_text fields: {st_text_updates}."
+ )
+
+
+if __name__ == "__main__":
+ main()
+
+# Example:
+# python normalize_jsonl_with_whisper.py \
+# --input /path/to/input.jsonl.gz \
+# --output /path/to/output.jsonl.gz
diff --git a/egs/europarl_st/SRT/local/normalize_texts.py b/egs/europarl_st/SRT/local/normalize_texts.py
new file mode 100644
index 0000000000..08464ee756
--- /dev/null
+++ b/egs/europarl_st/SRT/local/normalize_texts.py
@@ -0,0 +1,454 @@
+#!/usr/bin/env python3
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import logging
+import re
+import unicodedata
+from pathlib import Path
+from typing import Callable, Dict, Iterator, Optional, Sequence, Tuple
+
+try:
+ from whisper.normalizers import (
+ BasicTextNormalizer as WhisperBasicTextNormalizer, # type: ignore
+ )
+ from whisper.normalizers import (
+ EnglishTextNormalizer as WhisperEnglishTextNormalizer,
+ )
+except Exception: # pragma: no cover - only triggered when whisper isn't installed
+ WhisperBasicTextNormalizer = None
+ WhisperEnglishTextNormalizer = None
+
+
+class LocalBasicTextNormalizer:
+ """Language-agnostic text cleaner inspired by Whisper's BasicTextNormalizer."""
+
+ _WS_RE = re.compile(r"\s+")
+ _REPLACEMENTS = str.maketrans(
+ {
+ "\u2010": "-",
+ "\u2011": "-",
+ "\u2012": "-",
+ "\u2013": "-",
+ "\u2014": "-",
+ "\u2015": "-",
+ "\u2212": "-",
+ "\u2018": "'",
+ "\u2019": "'",
+ "\u201A": "'",
+ "\u201B": "'",
+ "\u2032": "'",
+ "\u2035": "'",
+ "\u201C": '"',
+ "\u201D": '"',
+ "\u201E": '"',
+ "\u00AB": '"',
+ "\u00BB": '"',
+ "\u02BC": "'",
+ "\u0060": "'",
+ "\u00B4": "'",
+ "\u200B": " ",
+ "\u200C": " ",
+ "\u200D": " ",
+ "\u200E": " ",
+ "\u200F": " ",
+ "\u202A": " ",
+ "\u202B": " ",
+ "\u202C": " ",
+ "\u202D": " ",
+ "\u202E": " ",
+ "\u2060": " ",
+ "\ufeff": " ",
+ "\u00A0": " ",
+ }
+ )
+
+ def __call__(self, text: str) -> str:
+ if text is None:
+ return ""
+ normalized = unicodedata.normalize("NFKC", str(text))
+ normalized = normalized.translate(self._REPLACEMENTS)
+ normalized = self._strip_control_characters(normalized)
+ normalized = self._WS_RE.sub(" ", normalized).strip()
+ return normalized
+
+ @staticmethod
+ def _strip_control_characters(text: str) -> str:
+ cleaned = []
+ for ch in text:
+ cat = unicodedata.category(ch)
+ if cat.startswith("C"):
+ # Preserve standard whitespace while collapsing it later.
+ cleaned.append(" " if ch.isspace() else "")
+ else:
+ cleaned.append(ch)
+ return "".join(cleaned)
+
+
+class LocalEnglishTextNormalizer(LocalBasicTextNormalizer):
+ """Very small subset of Whisper's EnglishTextNormalizer for offline use."""
+
+ _ASCII_ONLY_RE = re.compile(r"[^a-z0-9' ]+")
+
+ def __call__(self, text: str) -> str:
+ normalized = super().__call__(text)
+ normalized = normalized.lower()
+ normalized = self._ASCII_ONLY_RE.sub(" ", normalized)
+ normalized = self._WS_RE.sub(" ", normalized).strip()
+ return normalized
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(
+ description=(
+ "Apply Whisper-style text normalization to JSONL files and write "
+ "the results to a 'normalizer' subdirectory."
+ )
+ )
+ default_src = Path(__file__).resolve().parent.parent / "texts"
+ default_dst = default_src / "normalizer"
+ parser.add_argument(
+ "--src-dir",
+ type=Path,
+ default=default_src,
+ help="Root directory containing JSONL files (default: %(default)s).",
+ )
+ parser.add_argument(
+ "--dst-dir",
+ type=Path,
+ default=default_dst,
+ help="Destination root for normalized JSONL files (default: %(default)s).",
+ )
+ parser.add_argument(
+ "--fields",
+ nargs="+",
+ default=("text", "st_text"),
+ help=(
+ "JSON keys to normalize (default: text st_text). "
+ "Supports dot-separated paths."
+ ),
+ )
+ parser.add_argument(
+ "--normalizer",
+ choices=("basic", "english"),
+ default="basic",
+ help="Which normalization preset to use (default: basic).",
+ )
+ parser.add_argument(
+ "--skip-existing",
+ action="store_true",
+ help="Skip files whose normalized counterpart already exists.",
+ )
+ parser.add_argument(
+ "--dry-run",
+ action="store_true",
+ help="Do not write files; only report statistics.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Enable debug logging.",
+ )
+ parser.add_argument(
+ "--no-auto-english",
+ dest="auto_english",
+ action="store_false",
+ help=(
+ "Disable automatic detection of English fields based on "
+ "language-pair directory names."
+ ),
+ )
+ parser.set_defaults(auto_english=True)
+ return parser.parse_args()
+
+
+def load_normalizer(name: str):
+ if name == "english":
+ if WhisperEnglishTextNormalizer is not None:
+ logging.info("Using whisper.normalizers.EnglishTextNormalizer")
+ return WhisperEnglishTextNormalizer()
+ logging.info("Using built-in LocalEnglishTextNormalizer")
+ return LocalEnglishTextNormalizer()
+ if WhisperBasicTextNormalizer is not None:
+ logging.info("Using whisper.normalizers.BasicTextNormalizer")
+ return WhisperBasicTextNormalizer()
+ logging.info("Using built-in LocalBasicTextNormalizer")
+ return LocalBasicTextNormalizer()
+
+
+def iter_jsonl_files(src_dir: Path) -> Iterator[Path]:
+ for jsonl_path in sorted(src_dir.rglob("*.jsonl")):
+ if jsonl_path.is_file():
+ yield jsonl_path
+
+
+def ensure_destination_path(src_file: Path, src_root: Path, dst_root: Path) -> Path:
+ rel_path = src_file.relative_to(src_root)
+ dst_file = dst_root / rel_path
+ dst_file.parent.mkdir(parents=True, exist_ok=True)
+ return dst_file
+
+
+def set_nested_field(obj: dict, dotted_key: str, value: str) -> bool:
+ parts = dotted_key.split(".")
+ current = obj
+ for part in parts[:-1]:
+ if isinstance(current, dict) and part in current:
+ current = current[part]
+ else:
+ return False
+
+ last_key = parts[-1]
+ if isinstance(current, dict) and last_key in current:
+ current[last_key] = value
+ return True
+ return False
+
+
+def get_nested_field(obj: dict, dotted_key: str):
+ current = obj
+ for part in dotted_key.split("."):
+ if isinstance(current, dict):
+ if part not in current:
+ return None
+ current = current[part]
+ else:
+ return None
+ return current
+
+
+def normalize_line(
+ obj: dict,
+ fields: Sequence[str],
+ default_normalizer: Callable[[str], str],
+ field_normalizers: Dict[str, Callable[[str], str]],
+) -> int:
+ changed = 0
+ for field in fields:
+ text_value = get_nested_field(obj, field)
+ if not isinstance(text_value, str):
+ continue
+ normalize = field_normalizers.get(field, default_normalizer)
+ normalized = normalize(text_value).strip()
+ if normalized != text_value:
+ set_nested_field(obj, field, normalized)
+ changed += 1
+ return changed
+
+
+def process_file(
+ src_path: Path,
+ dst_path: Path,
+ fields: Sequence[str],
+ default_normalizer: Callable[[str], str],
+ field_normalizers: Dict[str, Callable[[str], str]],
+ dry_run: bool = False,
+) -> tuple[int, int]:
+ total = 0
+ changed_lines = 0
+ newline = "\n"
+ with src_path.open("r", encoding="utf-8") as src, (
+ open(dst_path, "w", encoding="utf-8", newline=newline)
+ if not dry_run
+ else nullcontext()
+ ) as dst: # type: ignore
+ for raw_line in src:
+ raw_line = raw_line.rstrip("\n")
+ if not raw_line:
+ continue
+ total += 1
+ obj = json.loads(raw_line)
+ changed = normalize_line(
+ obj,
+ fields,
+ default_normalizer,
+ field_normalizers,
+ )
+ if changed:
+ changed_lines += 1
+ if not dry_run:
+ dst.write(json.dumps(obj, ensure_ascii=False) + "\n")
+ return total, changed_lines
+
+
+class nullcontext:
+ """Minimal stand-in for contextlib.nullcontext for Python < 3.7 environments."""
+
+ def __enter__(self):
+ return None
+
+ def __exit__(self, exc_type, exc, tb):
+ return False
+
+
+def looks_like_lang_pair(name: str) -> bool:
+ return "_" in name and "." not in name
+
+
+def parse_lang_pair(name: str) -> Optional[Tuple[str, str]]:
+ if not looks_like_lang_pair(name):
+ return None
+ src, dst = name.split("_", 1)
+ if not src or not dst:
+ return None
+ return (src, dst)
+
+
+def detect_lang_pair(
+ src_dir: Path, src_file: Path
+) -> Tuple[Optional[str], Optional[str]]:
+ parsed = parse_lang_pair(src_dir.name)
+ if parsed:
+ return parsed
+ try:
+ relative = src_file.relative_to(src_dir)
+ except ValueError:
+ return (None, None)
+ if not relative.parts:
+ return (None, None)
+ parsed = parse_lang_pair(relative.parts[0])
+ if parsed:
+ return parsed
+ return (None, None)
+
+
+def choose_field_normalizers(
+ src_dir: Path,
+ src_file: Path,
+ fields: Sequence[str],
+ english_normalizer: Callable[[str], str],
+ auto_english: bool,
+) -> Dict[str, Callable[[str], str]]:
+ field_map: Dict[str, Callable[[str], str]] = {}
+ if not auto_english:
+ return field_map
+
+ src_lang, dst_lang = detect_lang_pair(src_dir, src_file)
+ if not src_lang and not dst_lang:
+ return field_map
+
+ src_lang = (src_lang or "").lower()
+ dst_lang = (dst_lang or "").lower()
+
+ if src_lang.startswith("en") and "text" in fields:
+ field_map["text"] = english_normalizer
+ if dst_lang.startswith("en") and "st_text" in fields:
+ field_map["st_text"] = english_normalizer
+ return field_map
+
+
+def main():
+ args = parse_args()
+ logging.basicConfig(
+ format="%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s",
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ )
+
+ basic_normalizer = load_normalizer("basic")
+ english_normalizer = load_normalizer("english")
+ default_normalizer = (
+ english_normalizer if args.normalizer == "english" else basic_normalizer
+ )
+ args.src_dir = args.src_dir.resolve()
+ args.dst_dir = args.dst_dir.resolve()
+ args.dst_dir.mkdir(parents=True, exist_ok=True)
+
+ logging.info(
+ "Normalizing JSONL files from %s into %s (fields=%s, dry_run=%s)",
+ args.src_dir,
+ args.dst_dir,
+ ", ".join(args.fields),
+ args.dry_run,
+ )
+
+ processed_files = 0
+ total_lines = 0
+ total_changed = 0
+
+ for src_file in iter_jsonl_files(args.src_dir):
+ dst_file = ensure_destination_path(src_file, args.src_dir, args.dst_dir)
+ if args.skip_existing and dst_file.is_file():
+ logging.info("Skipping %s (already exists)", src_file)
+ continue
+
+ logging.info("Processing %s -> %s", src_file, dst_file)
+ field_normalizers = choose_field_normalizers(
+ args.src_dir,
+ src_file,
+ args.fields,
+ english_normalizer,
+ args.auto_english,
+ )
+ if field_normalizers and args.verbose:
+ logging.debug(
+ "English normalizer applied to fields %s for %s",
+ ", ".join(sorted(field_normalizers)),
+ src_file,
+ )
+ file_total, file_changed = process_file(
+ src_file,
+ dst_file,
+ args.fields,
+ default_normalizer,
+ field_normalizers,
+ dry_run=args.dry_run,
+ )
+ processed_files += 1
+ total_lines += file_total
+ total_changed += file_changed
+ logging.info(
+ "Finished %s (lines=%d, changed=%d)",
+ src_file.name,
+ file_total,
+ file_changed,
+ )
+
+ logging.info(
+ "Done. files=%d, lines=%d, lines_changed=%d",
+ processed_files,
+ total_lines,
+ total_changed,
+ )
+
+
+if __name__ == "__main__":
+ main()
+
+
+"""
+Example usage:
+
+python normalize_texts.py \
+ --src-dir ./texts \
+ --dst-dir ./normalizer \
+ --fields text st_text \
+ --normalizer basic \
+ --skip-existing
+
+Options:
+ --src-dir / --dst-dir: Input and output directories (defaults are relative to this script).
+ --fields: JSON keys to normalize (supports dot-separated paths like nested.field).
+ Default: text st_text.
+ --normalizer basic|english: Choose the base normalizer. If whisper.normalizers is installed,
+ the official implementation is used; otherwise falls back to built-in
+ LocalBasicTextNormalizer / LocalEnglishTextNormalizer.
+ --auto-english (enabled by default): Automatically detects English fields from directory
+ names like 'es_en' and applies the English normalizer (keeps only [a-z0-9']).
+ Use --no-auto-english to disable.
+ --skip-existing: Skip files that already have a corresponding output.
+ --dry-run: Only report statistics without writing to disk.
+ --verbose: Enable DEBUG-level logging.
+"""
diff --git a/egs/europarl_st/SRT/local/org_to_jsonl.py b/egs/europarl_st/SRT/local/org_to_jsonl.py
new file mode 100644
index 0000000000..4f701535cf
--- /dev/null
+++ b/egs/europarl_st/SRT/local/org_to_jsonl.py
@@ -0,0 +1,231 @@
+# Copyright 2026 Nanjie Li (linanjie0820@gmail.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Preprocess the Europarl-ST dataset into per-language-pair JSONL files.
+
+Reference: https://www.mllp.upv.es/europarl-st/
+"""
+
+import argparse
+import json
+import os
+import sys
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
+
+from utils.audio_utils import audio_to_flac
+from utils.dataset_parameters import AUDIO_SAVE_SAMPLE_RATE
+
+
+def read_lst_file(filename):
+ with open(filename) as f:
+ lines = [line.rstrip() for line in f]
+ return lines
+
+
+def parse_timestamp(timestamp):
+ if ":" in timestamp:
+ parts = [float(part) for part in timestamp.split(":")]
+ while len(parts) < 3:
+ parts.insert(0, 0.0)
+ hours, minutes, seconds = parts[-3], parts[-2], parts[-1]
+ return hours * 3600 + minutes * 60 + seconds
+ return float(timestamp)
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(
+ description="Convert raw Europarl-ST data to per-language-pair JSONL files."
+ )
+ parser.add_argument(
+ "--data-dir",
+ type=str,
+ default=os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "v1.1"),
+ help="Path to the Europarl-ST v1.1 raw data directory.",
+ )
+ parser.add_argument(
+ "--output-dir",
+ type=str,
+ default=os.path.join(
+ os.path.dirname(os.path.realpath(__file__)), "..", "audio"
+ ),
+ help="Directory to store converted FLAC audio segments.",
+ )
+ return parser.parse_args()
+
+
+if __name__ == "__main__":
+ args = parse_args()
+ data_dir = os.path.realpath(args.data_dir)
+ output_dir = os.path.realpath(args.output_dir)
+
+ # Remap original splits: train/dev/test -> train/valid/test.
+ split_output_name_dict = {
+ "train": "train",
+ "dev": "valid",
+ "test": "test",
+ # 'train-noisy': 'train-noisy' # Skipped: this split contains many errors.
+ }
+
+ languages = ["es", "de", "en", "fr", "nl", "pl", "pt", "ro", "it"]
+
+ texts_output_dir = os.path.normpath(
+ os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "texts")
+ )
+ os.makedirs(texts_output_dir, exist_ok=True)
+
+ pair_jsonl_paths = {}
+ dataset_prefix = "europarl"
+
+ for src_lang in languages:
+ for dest_lang in languages:
+ if src_lang == dest_lang:
+ continue
+ pair_dir = os.path.join(texts_output_dir, f"{src_lang}_{dest_lang}")
+ os.makedirs(pair_dir, exist_ok=True)
+ for split_alias in split_output_name_dict.values():
+ jsonl_path = os.path.join(
+ pair_dir,
+ f"{dataset_prefix}.{src_lang}_{dest_lang}.{split_alias}.jsonl",
+ )
+ with open(jsonl_path, "w", encoding="utf-8"):
+ pass
+ pair_jsonl_paths[(src_lang, dest_lang, split_alias)] = jsonl_path
+
+ dict_skeleton = {
+ "es": None,
+ "de": None,
+ "en": None,
+ "fr": None,
+ "nl": None,
+ "pl": None,
+ "pt": None,
+ "ro": None,
+ "it": None,
+ }
+
+ file_ids = {"train": 1, "dev": 1, "test": 1}
+
+ for source_lang in languages:
+
+ print(f"Processing {source_lang} dataset...\n")
+
+ language_folder = os.path.join(data_dir, source_lang)
+
+ destination_languages = languages.copy()
+ destination_languages.remove(source_lang)
+
+ for split, split_name in split_output_name_dict.items():
+
+ os.makedirs(os.path.join(output_dir, split_name), exist_ok=True)
+ segments_dict = {}
+
+ for dest_lang in destination_languages:
+
+ segments_lst_file = os.path.join(
+ language_folder, dest_lang, split, "segments.lst"
+ )
+ segments_source_lang_file = os.path.join(
+ language_folder, dest_lang, split, f"segments.{source_lang}"
+ )
+ segments_dest_lang_file = os.path.join(
+ language_folder, dest_lang, split, f"segments.{dest_lang}"
+ )
+
+ segments_timestamps = read_lst_file(segments_lst_file)
+ segments_source_lang_transcriptions = read_lst_file(
+ segments_source_lang_file
+ )
+ segments_dest_lang_transcriptions = read_lst_file(
+ segments_dest_lang_file
+ )
+
+ segments_source_lang_transcriptions_dict = dict(
+ zip(segments_timestamps, segments_source_lang_transcriptions)
+ )
+ segments_dest_lang_transcriptions_dict = dict(
+ zip(segments_timestamps, segments_dest_lang_transcriptions)
+ )
+
+ for segment in segments_timestamps:
+
+ segments_dict.setdefault(segment, dict_skeleton.copy())
+ segments_dict[segment][
+ source_lang
+ ] = segments_source_lang_transcriptions_dict[segment]
+ segments_dict[segment][
+ dest_lang
+ ] = segments_dest_lang_transcriptions_dict[segment]
+
+ for segment in segments_dict:
+ audio, segment_start, segment_end = segment.split()
+
+ audio_path = os.path.join(language_folder, "audios", f"{audio}.m4a")
+ audio_output_filename = f"{source_lang}_{file_ids[split]}.flac"
+ audio_output_path = os.path.join(
+ output_dir, split_name, audio_output_filename
+ )
+
+ audio_to_flac(
+ audio_path,
+ audio_output_path,
+ sample_rate=AUDIO_SAVE_SAMPLE_RATE,
+ segment_start=segment_start,
+ segment_end=segment_end,
+ )
+
+ source_transcription = segments_dict[segment][source_lang]
+ if source_transcription in (None, "None"):
+ file_ids[split] += 1
+ continue
+
+ translations = {
+ lang: text
+ for lang, text in segments_dict[segment].items()
+ if lang != source_lang and text not in (None, "None")
+ }
+
+ if not translations:
+ file_ids[split] += 1
+ continue
+
+ try:
+ duration_seconds = parse_timestamp(segment_end) - parse_timestamp(
+ segment_start
+ )
+ except ValueError:
+ file_ids[split] += 1
+ continue
+
+ if duration_seconds <= 0:
+ file_ids[split] += 1
+ continue
+
+ rounded_duration = round(duration_seconds, 3)
+
+ for dest_lang, translation in translations.items():
+ jsonl_path = pair_jsonl_paths[(source_lang, dest_lang, split_name)]
+ jsonl_entry = {
+ "source": audio_output_path,
+ "duration": rounded_duration,
+ "text": source_transcription,
+ "st_text": translation,
+ }
+ with open(jsonl_path, "a", encoding="utf-8") as jsonl_file:
+ jsonl_file.write(
+ json.dumps(jsonl_entry, ensure_ascii=False) + "\n"
+ )
+
+ file_ids[split] += 1
diff --git a/egs/europarl_st/SRT/local/texts_to_cuts.py b/egs/europarl_st/SRT/local/texts_to_cuts.py
new file mode 100644
index 0000000000..ca3693660b
--- /dev/null
+++ b/egs/europarl_st/SRT/local/texts_to_cuts.py
@@ -0,0 +1,460 @@
+#!/usr/bin/env python3
+# Copyright 2025 Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import logging
+from dataclasses import replace
+from itertools import chain
+from pathlib import Path
+from typing import Dict, Iterator, List, Tuple
+
+try:
+ import torch # type: ignore
+except ModuleNotFoundError as exc: # pragma: no cover
+ raise RuntimeError(
+ "texts_to_cuts.py now requires torch. Please install PyTorch."
+ ) from exc
+
+try:
+ from lhotse import (
+ CutSet,
+ Features,
+ KaldifeatFbank,
+ KaldifeatFbankConfig,
+ LilcomChunkyWriter,
+ Recording,
+ RecordingSet,
+ SupervisionSegment,
+ SupervisionSet,
+ )
+except ModuleNotFoundError as exc: # pragma: no cover
+ raise RuntimeError(
+ "texts_to_cuts.py now depends on lhotse (pip install lhotse)."
+ ) from exc
+
+
+def parse_args() -> argparse.Namespace:
+ default_src = Path(__file__).resolve().parent.parent / "texts" / "normalizer"
+ default_dst = (default_src.parents[1] / "cut_manifests").resolve()
+ default_storage = (default_dst.parents[0] / "fbank_storage").resolve()
+ parser = argparse.ArgumentParser(
+ description=(
+ "Convert normalized JSONL files into CutSet manifests that already "
+ "contain kaldifeat FBANK features (MonoCut entries with 'features')."
+ )
+ )
+ parser.add_argument(
+ "--src-dir",
+ type=Path,
+ default=default_src,
+ help="Root directory containing normalized *.jsonl files.",
+ )
+ parser.add_argument(
+ "--dst-dir",
+ type=Path,
+ default=default_dst,
+ help="Output directory for generated *.jsonl.gz files (mirrors src layout).",
+ )
+ parser.add_argument(
+ "--audio-root",
+ type=Path,
+ default=Path(__file__).resolve().parent.parent.parent,
+ help="Base directory that the JSON 'source' field is relative to.",
+ )
+ parser.add_argument(
+ "--storage-root",
+ type=Path,
+ default=default_storage,
+ help="Directory where .lca feature chunks will be written.",
+ )
+ parser.add_argument(
+ "--num-workers",
+ type=int,
+ default=8,
+ help="Number of workers used when computing FBANK features.",
+ )
+ parser.add_argument(
+ "--batch-duration",
+ type=float,
+ default=600.0,
+ help="Total audio seconds per minibatch for feature extraction.",
+ )
+ parser.add_argument(
+ "--device",
+ type=str,
+ default="auto",
+ help="Device for feature extraction (auto/cpu/cuda).",
+ )
+ parser.add_argument(
+ "--train-shard-duration",
+ type=float,
+ default=0.0,
+ help=(
+ "For train splits only: limit total audio seconds per feature shard. "
+ "Each shard writes its own .lca when > 0 (disabled by default)."
+ ),
+ )
+ parser.add_argument(
+ "--skip-missing-audio",
+ action="store_true",
+ help="Skip entries whose audio file is missing instead of raising.",
+ )
+ parser.add_argument(
+ "--feature-cache",
+ type=Path,
+ default=default_storage / "feature_cache.json",
+ help="JSON file used to memoize feature descriptors per recording.",
+ )
+ parser.add_argument(
+ "--refresh-cache",
+ action="store_true",
+ help="Ignore any existing cache entries and recompute features.",
+ )
+ parser.add_argument(
+ "--overwrite",
+ action="store_true",
+ help="Overwrite existing *.jsonl.gz outputs.",
+ )
+ parser.add_argument(
+ "--verbose",
+ action="store_true",
+ help="Enable debug logging.",
+ )
+ return parser.parse_args()
+
+
+def iter_jsonl_files(root: Path) -> Iterator[Path]:
+ for path in sorted(root.rglob("*.jsonl")):
+ if path.is_file():
+ yield path
+
+
+def ensure_destination_path(src_file: Path, src_root: Path, dst_root: Path) -> Path:
+ rel = src_file.relative_to(src_root)
+ base = src_file.name[:-6] if src_file.name.endswith(".jsonl") else src_file.stem
+ dst_name = f"{base}_cuts.jsonl.gz"
+ dst_path = dst_root / rel.parent / dst_name
+ dst_path.parent.mkdir(parents=True, exist_ok=True)
+ return dst_path
+
+
+def parse_lang_pair(path: Path) -> Tuple[str, str]:
+ parent_name = path.parent.name
+ if "_" not in parent_name:
+ return ("unknown", "unknown")
+ src, tgt = parent_name.split("_", 1)
+ return (src, tgt)
+
+
+def resolve_audio_path(audio_root: Path, raw_source: str) -> Path:
+ raw_path = Path(raw_source)
+ if not raw_path.is_absolute():
+ raw_path = (audio_root / raw_source).resolve()
+ return raw_path
+
+
+def build_manifest_entries(
+ src_file: Path,
+ audio_root: Path,
+ skip_missing_audio: bool,
+) -> Tuple[RecordingSet, SupervisionSet, int, int]:
+ src_lang, tgt_lang = parse_lang_pair(src_file)
+ recordings: Dict[str, Recording] = {}
+ supervisions = []
+ total = 0
+ kept = 0
+ with src_file.open("r", encoding="utf-8") as inp:
+ for raw_line in inp:
+ raw_line = raw_line.strip()
+ if not raw_line:
+ continue
+ total += 1
+ item = json.loads(raw_line)
+ audio_path = resolve_audio_path(audio_root, item["source"])
+ if not audio_path.is_file():
+ msg = f"Missing audio: {audio_path}"
+ if skip_missing_audio:
+ logging.warning("%s (skipping)", msg)
+ continue
+ raise FileNotFoundError(msg)
+
+ recording_id = audio_path.stem
+ if recording_id not in recordings:
+ recordings[recording_id] = Recording.from_file(
+ path=str(audio_path), recording_id=recording_id
+ )
+
+ supervision = SupervisionSegment(
+ id=recording_id,
+ recording_id=recording_id,
+ start=0.0,
+ duration=recordings[recording_id].duration,
+ channel=0,
+ text=item.get("text", ""),
+ language=src_lang,
+ speaker="unknown",
+ gender="",
+ custom={
+ "st_text": item.get("st_text", ""),
+ "lang": tgt_lang,
+ },
+ )
+ supervisions.append(supervision)
+ kept += 1
+
+ recording_set = RecordingSet.from_recordings(recordings.values())
+ supervision_set = SupervisionSet.from_segments(supervisions)
+ return recording_set, supervision_set, total, kept
+
+
+def split_cached_cuts(
+ cut_set: CutSet,
+ feature_cache: Dict[str, Dict],
+ refresh_cache: bool,
+) -> Tuple[List, List]:
+ cached_cuts = []
+ pending_cuts = []
+ for cut in cut_set:
+ rid = cut.recording.id
+ if not refresh_cache and rid in feature_cache:
+ feat = Features.from_dict(feature_cache[rid])
+ cached_cuts.append(replace(cut, features=feat))
+ else:
+ pending_cuts.append(cut)
+ return cached_cuts, pending_cuts
+
+
+def update_feature_cache(feature_cache_path: Path, cache: Dict[str, Dict]) -> None:
+ feature_cache_path.parent.mkdir(parents=True, exist_ok=True)
+ with feature_cache_path.open("w", encoding="utf-8") as f:
+ json.dump(cache, f)
+
+
+def load_feature_cache(feature_cache_path: Path) -> Dict[str, Dict]:
+ if not feature_cache_path.is_file():
+ return {}
+ with feature_cache_path.open("r", encoding="utf-8") as f:
+ return json.load(f)
+
+
+def shard_cuts_by_duration(cuts: List, max_duration: float) -> List[List]:
+ """Group cuts into shards whose total duration does not exceed max_duration."""
+ if max_duration <= 0:
+ return [cuts]
+
+ shards: List[List] = []
+ current: List = []
+ current_dur = 0.0
+
+ for cut in cuts:
+ dur = float(cut.duration or 0.0)
+ if current and current_dur + dur > max_duration:
+ shards.append(current)
+ current = []
+ current_dur = 0.0
+ current.append(cut)
+ current_dur += dur
+
+ if current:
+ shards.append(current)
+
+ return [shard for shard in shards if shard]
+
+
+def add_dataloading_info(cut):
+ info = {"rank": 0, "world_size": 1, "worker_id": None}
+ return cut.with_custom("dataloading_info", info)
+
+
+def process_file(
+ src_file: Path,
+ dst_file: Path,
+ audio_root: Path,
+ storage_root: Path,
+ feature_cache: Dict[str, Dict],
+ feature_cache_path: Path,
+ num_workers: int,
+ batch_duration: float,
+ device: str,
+ train_shard_duration: float,
+ skip_missing_audio: bool,
+ refresh_cache: bool,
+) -> Tuple[int, int]:
+ storage_root.mkdir(parents=True, exist_ok=True)
+ recordings, supervisions, total, kept = build_manifest_entries(
+ src_file=src_file,
+ audio_root=audio_root,
+ skip_missing_audio=skip_missing_audio,
+ )
+ if kept == 0:
+ logging.warning("No usable entries in %s; skipping.", src_file)
+ return total, kept
+
+ cut_set = CutSet.from_manifests(
+ recordings=recordings,
+ supervisions=supervisions,
+ )
+
+ cached_cuts, pending_cuts = split_cached_cuts(
+ cut_set=cut_set,
+ feature_cache=feature_cache,
+ refresh_cache=refresh_cache,
+ )
+
+ rel = src_file.parent.relative_to(src_file.parents[1])
+ base = src_file.stem
+ storage_path = storage_root / rel / base
+ storage_path.parent.mkdir(parents=True, exist_ok=True)
+
+ actual_device = device
+ if device == "auto":
+ actual_device = "cuda" if torch.cuda.is_available() else "cpu"
+ extractor = KaldifeatFbank(KaldifeatFbankConfig(device=actual_device))
+
+ new_cuts = pending_cuts
+ should_shard = train_shard_duration > 0 and "train" in src_file.stem and new_cuts
+
+ if new_cuts:
+ if should_shard:
+ shards = shard_cuts_by_duration(new_cuts, train_shard_duration)
+ else:
+ shards = [new_cuts]
+
+ shard_cut_sets = []
+ for idx, shard in enumerate(shards):
+ shard_storage = (
+ f"{storage_path}_shard{idx:04d}"
+ if len(shards) > 1
+ else str(storage_path)
+ )
+ shard_set = CutSet.from_cuts(shard)
+ shard_set = shard_set.compute_and_store_features_batch(
+ extractor=extractor,
+ storage_path=shard_storage,
+ storage_type=LilcomChunkyWriter,
+ num_workers=num_workers,
+ batch_duration=batch_duration,
+ overwrite=True,
+ )
+ for cut in shard_set:
+ feature_cache[cut.recording.id] = cut.features.to_dict()
+ shard_cut_sets.append(shard_set)
+
+ new_cut_set = CutSet.from_cuts(chain.from_iterable(shard_cut_sets))
+ cached_cuts.extend(new_cut_set)
+ update_feature_cache(feature_cache_path, feature_cache)
+
+ final_cut_set = CutSet.from_cuts(cached_cuts)
+
+ final_cut_set = final_cut_set.map(add_dataloading_info)
+ final_cut_set.to_file(dst_file)
+ return total, kept
+
+
+def main():
+ args = parse_args()
+ logging.basicConfig(
+ format="%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s",
+ level=logging.DEBUG if args.verbose else logging.INFO,
+ )
+ args.src_dir = args.src_dir.resolve()
+ args.dst_dir = args.dst_dir.resolve()
+ args.dst_dir.mkdir(parents=True, exist_ok=True)
+ audio_root = args.audio_root.resolve()
+ storage_root = args.storage_root.resolve()
+ feature_cache_path = args.feature_cache.resolve()
+ feature_cache = {} if args.refresh_cache else load_feature_cache(feature_cache_path)
+
+ logging.info(
+ "Converting normalized texts from %s into feature-rich CutSet manifests under %s",
+ args.src_dir,
+ args.dst_dir,
+ )
+
+ processed_files = 0
+ total_lines = 0
+ total_kept = 0
+ for src_file in iter_jsonl_files(args.src_dir):
+ dst_file = ensure_destination_path(src_file, args.src_dir, args.dst_dir)
+ if dst_file.is_file() and not args.overwrite:
+ logging.info("Skipping %s (exists)", dst_file)
+ continue
+ logging.info("Processing %s -> %s", src_file, dst_file)
+ file_total, file_kept = process_file(
+ src_file=src_file,
+ dst_file=dst_file,
+ audio_root=audio_root,
+ storage_root=storage_root,
+ feature_cache=feature_cache,
+ feature_cache_path=feature_cache_path,
+ num_workers=args.num_workers,
+ batch_duration=args.batch_duration,
+ device=args.device,
+ train_shard_duration=args.train_shard_duration,
+ skip_missing_audio=args.skip_missing_audio,
+ refresh_cache=args.refresh_cache,
+ )
+ processed_files += 1
+ total_lines += file_total
+ total_kept += file_kept
+ logging.info(
+ "Finished %s (lines=%d, kept=%d)",
+ src_file.name,
+ file_total,
+ file_kept,
+ )
+
+ logging.info(
+ "Done. files=%d, lines=%d, kept=%d",
+ processed_files,
+ total_lines,
+ total_kept,
+ )
+
+
+if __name__ == "__main__":
+ main()
+
+
+"""
+Example usage:
+
+python texts_to_cuts.py \
+ --src-dir ./normalizer \
+ --dst-dir ./manifests \
+ --audio-root /path/to/audio/root \
+ --storage-root ./fbank \
+ --feature-cache ./fbank/feature_cache.json \
+ --num-workers 12 \
+ --batch-duration 400 \
+ --skip-missing-audio \
+ --overwrite \
+ --verbose
+
+Options:
+ --src-dir: Root directory containing normalized *.jsonl files (recursively traversed).
+ --dst-dir: Output root for *_cuts.jsonl.gz manifests (mirrors the src layout).
+ --audio-root: Base path that the JSON 'source' field is relative to.
+ --storage-root: Root directory for storing FBANK .lca feature chunks.
+ --feature-cache: JSON file for memoizing extracted features; reuses cached features
+ for previously seen audio files to avoid redundant computation.
+ --refresh-cache: If specified, ignores the cache and recomputes all features.
+ --num-workers / --batch-duration / --device: Control parallelism and device for
+ feature extraction (auto prefers GPU if available).
+ --skip-missing-audio: Skip entries with missing audio files instead of raising an error.
+ --overwrite: Overwrite existing *_cuts.jsonl.gz outputs.
+ --verbose: Enable more detailed logging.
+"""
diff --git a/egs/europarl_st/SRT/local/train_bpe.py b/egs/europarl_st/SRT/local/train_bpe.py
new file mode 100755
index 0000000000..72378ccdb3
--- /dev/null
+++ b/egs/europarl_st/SRT/local/train_bpe.py
@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# You can install sentencepiece via:
+#
+# pip install sentencepiece
+#
+# Due to an issue reported in
+# https://github.com/google/sentencepiece/pull/642#issuecomment-857972030
+#
+# Please install a version >=0.1.96
+
+import argparse
+import shutil
+from pathlib import Path
+from typing import Dict
+
+import sentencepiece as spm
+
+
+def get_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--lang-dir",
+ type=str,
+ help="""Input and output directory.
+ The generated bpe.model is saved to this directory.
+ """,
+ )
+
+ parser.add_argument(
+ "--transcript",
+ type=str,
+ help="Training transcript.",
+ )
+
+ parser.add_argument(
+ "--vocab-size",
+ type=int,
+ help="Vocabulary size for BPE training",
+ )
+
+ return parser.parse_args()
+
+
+def generate_tokens(lang_dir: Path):
+ """
+ Generate the tokens.txt from a bpe model.
+ """
+ sp = spm.SentencePieceProcessor()
+ sp.load(str(lang_dir / "bpe.model"))
+ token2id: Dict[str, int] = {sp.id_to_piece(i): i for i in range(sp.vocab_size())}
+ with open(lang_dir / "tokens.txt", "w", encoding="utf-8") as f:
+ for sym, i in token2id.items():
+ f.write(f"{sym} {i}\n")
+
+
+def main():
+ args = get_args()
+ vocab_size = args.vocab_size
+ lang_dir = Path(args.lang_dir)
+
+ model_type = "unigram"
+
+ model_prefix = f"{lang_dir}/{model_type}_{vocab_size}"
+ train_text = args.transcript
+ character_coverage = 1.0
+ input_sentence_size = 100000000
+
+ user_defined_symbols = [
+ "",
+ "",
+ "<2en>",
+ "<2de>",
+ "<2es>",
+ "<2fr>",
+ "<2it>",
+ "<2nl>",
+ "<2pl>",
+ "<2pt>",
+ "<2ro>",
+ ]
+
+ unk_id = len(user_defined_symbols)
+
+ model_file = Path(model_prefix + ".model")
+ if not model_file.is_file():
+ spm.SentencePieceTrainer.train(
+ input=train_text,
+ vocab_size=vocab_size,
+ model_type=model_type,
+ model_prefix=model_prefix,
+ shuffle_input_sentence=True,
+ input_sentence_size=input_sentence_size,
+ character_coverage=character_coverage,
+ user_defined_symbols=user_defined_symbols,
+ unk_id=unk_id,
+ bos_id=-1,
+ eos_id=-1,
+ train_extremely_large_corpus=True,
+ )
+ else:
+ print(f"{model_file} exists - skipping")
+ return
+
+ shutil.copyfile(model_file, f"{lang_dir}/bpe.model")
+
+ generate_tokens(lang_dir)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/egs/europarl_st/SRT/local/utils/__init__.py b/egs/europarl_st/SRT/local/utils/__init__.py
new file mode 100644
index 0000000000..13ee1d163d
--- /dev/null
+++ b/egs/europarl_st/SRT/local/utils/__init__.py
@@ -0,0 +1,3 @@
+"""Utility package for Europarl-ST preprocessing scripts."""
+
+__all__ = ["audio_utils", "dataset_parameters"]
diff --git a/egs/europarl_st/SRT/local/utils/audio_utils.py b/egs/europarl_st/SRT/local/utils/audio_utils.py
new file mode 100644
index 0000000000..90794109fe
--- /dev/null
+++ b/egs/europarl_st/SRT/local/utils/audio_utils.py
@@ -0,0 +1,86 @@
+"""
+Audio conversion helpers for Europarl-ST preprocessing.
+
+Currently exposes `audio_to_flac`, a thin wrapper around FFmpeg that
+extracts (and optionally trims) segments from the original recordings
+while resampling to the desired rate.
+"""
+
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+from typing import Optional
+
+
+class AudioConversionError(RuntimeError):
+ """Raised when FFmpeg cannot convert the provided audio snippet."""
+
+
+def _ensure_parent_dir(path: Path) -> None:
+ """Create the parent directory for `path` if it does not exist."""
+ path.parent.mkdir(parents=True, exist_ok=True)
+
+
+def audio_to_flac(
+ input_path: os.PathLike[str] | str,
+ output_path: os.PathLike[str] | str,
+ sample_rate: int,
+ segment_start: Optional[str] = None,
+ segment_end: Optional[str] = None,
+) -> None:
+ """
+ Convert `input_path` audio into a FLAC file at `output_path`.
+
+ Args:
+ input_path: Source audio file (.m4a in the Europarl-ST dataset).
+ output_path: Destination path for the trimmed/resampled FLAC.
+ sample_rate: Target sample rate in Hz (e.g., 16000).
+ segment_start: Optional HH:MM:SS.sss start timestamp.
+ segment_end: Optional HH:MM:SS.sss end timestamp.
+
+ Raises:
+ FileNotFoundError: If the source audio does not exist.
+ AudioConversionError: When FFmpeg fails.
+ """
+
+ input_path = Path(input_path)
+ output_path = Path(output_path)
+
+ if not input_path.is_file():
+ raise FileNotFoundError(f"Audio source not found: {input_path}")
+
+ _ensure_parent_dir(output_path)
+
+ cmd = ["ffmpeg", "-y", "-hide_banner", "-loglevel", "error", "-i", str(input_path)]
+
+ if segment_start is not None:
+ cmd += ["-ss", str(segment_start)]
+
+ if segment_end is not None:
+ cmd += ["-to", str(segment_end)]
+
+ cmd += [
+ "-ar",
+ str(int(sample_rate)),
+ "-ac",
+ "1",
+ "-c:a",
+ "flac",
+ str(output_path),
+ ]
+
+ proc = subprocess.run(
+ cmd,
+ stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE,
+ check=False,
+ encoding="utf-8",
+ errors="ignore",
+ )
+
+ if proc.returncode != 0:
+ raise AudioConversionError(
+ f"FFmpeg failed ({proc.returncode}) while converting {input_path} -> {output_path}:\n{proc.stderr}"
+ )
diff --git a/egs/europarl_st/SRT/local/utils/dataset_parameters.py b/egs/europarl_st/SRT/local/utils/dataset_parameters.py
new file mode 100644
index 0000000000..0e72d5d62c
--- /dev/null
+++ b/egs/europarl_st/SRT/local/utils/dataset_parameters.py
@@ -0,0 +1,5 @@
+"""
+Shared constants for Europarl-ST preprocessing.
+"""
+
+AUDIO_SAVE_SAMPLE_RATE = 16000
diff --git a/egs/europarl_st/SRT/prepare.sh b/egs/europarl_st/SRT/prepare.sh
new file mode 100755
index 0000000000..31635e617e
--- /dev/null
+++ b/egs/europarl_st/SRT/prepare.sh
@@ -0,0 +1,96 @@
+#!/usr/bin/env bash
+
+set -eou pipefail
+
+# This script prepares the Europarl-ST dataset for joint multilingual
+# speech recognition and translation (SRT) training.
+#
+# Usage:
+# cd egs/europarl_st/SRT
+# bash prepare.sh --stage 0 --stop-stage 5
+#
+# Prerequisites:
+# - Download Europarl-ST v1.1 from https://www.mllp.upv.es/europarl-st/
+# - Place it under data/europarl_st/v1.1/
+
+stage=0
+stop_stage=5
+
+# Paths (modify as needed)
+data_root=./data/europarl_st
+raw_dir=${data_root}/v1.1
+audio_dir=${data_root}/audio
+texts_dir=${data_root}/texts
+norm_dir=${data_root}/normalizer
+manifest_dir=${data_root}/manifests
+fbank_dir=${data_root}/fbank
+bpe_dir=${data_root}/bpe
+
+. shared/parse_options.sh || exit 1
+
+log() {
+ local fname=${BASH_SOURCE[1]##*/}
+ echo -e "$(date '+%Y-%m-%dT%H:%M:%S') (${fname}:${BASH_LINENO[0]}:${FUNCNAME[1]}) $*"
+}
+
+if [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then
+ log "Stage 0: Extract audio segments and build JSONL"
+ python local/org_to_jsonl.py \
+ --data-dir ${raw_dir} \
+ --output-dir ${audio_dir}
+fi
+
+if [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then
+ log "Stage 1: Normalize text"
+ python local/normalize_texts.py \
+ --src-dir ${texts_dir} \
+ --dst-dir ${norm_dir} \
+ --fields text st_text \
+ --normalizer basic \
+ --skip-existing
+fi
+
+if [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then
+ log "Stage 2: Generate CutSet manifests with FBANK features"
+ python local/texts_to_cuts.py \
+ --src-dir ${norm_dir} \
+ --dst-dir ${manifest_dir} \
+ --audio-root ${data_root} \
+ --storage-root ${fbank_dir} \
+ --num-workers 8 \
+ --skip-missing-audio
+fi
+
+if [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then
+ log "Stage 3: Filter out entries with empty text"
+ python local/filter_cuts_texts.py \
+ --manifest-dir ${manifest_dir} \
+ --overwrite
+fi
+
+if [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then
+ log "Stage 4: Validate manifests"
+ python local/check_manifests.py \
+ --manifests-dir ${manifest_dir} \
+ --num-workers 8
+fi
+
+if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
+ log "Stage 5: Train BPE models"
+
+ # ASR BPE (9-language shared)
+ mkdir -p ${bpe_dir}/asr9
+ python local/train_bpe.py \
+ --lang-dir ${bpe_dir}/asr9 \
+ --transcript ${data_root}/asr_train_text.txt \
+ --vocab-size 500
+
+ # ST BPE (9-language shared)
+ mkdir -p ${bpe_dir}/ast9
+ python local/train_bpe.py \
+ --lang-dir ${bpe_dir}/ast9 \
+ --transcript ${data_root}/ast_train_text.txt \
+ --vocab-size 6000
+fi
+
+log "Data preparation completed successfully!"
diff --git a/egs/europarl_st/SRT/shared b/egs/europarl_st/SRT/shared
new file mode 120000
index 0000000000..4c5e91438c
--- /dev/null
+++ b/egs/europarl_st/SRT/shared
@@ -0,0 +1 @@
+../../../icefall/shared/
\ No newline at end of file