Merge branch 'k2-fsa:master' into dev_zipformer_cn

2025-09-04 14:44:18 +00:00 · 2023-10-23 09:11:14 +08:00 · 2023-10-23 09:11:14 +08:00 · 3eb55080aa
commit 3eb55080aa
parent 023f6e05d4 902dc2364a
359 changed files with 44997 additions and 843 deletions
--- a/.flake8
+++ b/.flake8
@ -24,6 +24,7 @@ exclude =
  **/data/**,
  icefall/shared/make_kn_lm.py,
  icefall/__init__.py
+  icefall/ctc/__init__.py

 ignore =
  # E203 white space before ":"
--- a/.github/scripts/run-gigaspeech-zipformer-2023-10-17.sh
+++ b/.github/scripts/run-gigaspeech-zipformer-2023-10-17.sh
@ -0,0 +1,94 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+cd egs/gigaspeech/ASR
+
+repo_url=https://huggingface.co/yfyeung/icefall-asr-gigaspeech-zipformer-2023-10-17
+
+log "Downloading pre-trained model from $repo_url"
+git lfs install
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+
+log "Display test files"
+tree $repo/
+ls -lh $repo/test_wavs/*.wav
+
+pushd $repo/exp
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "data/lang_bpe_500/tokens.txt"
+git lfs pull --include "exp/jit_script.pt"
+git lfs pull --include "exp/pretrained.pt"
+ln -s pretrained.pt epoch-99.pt
+ls -lh *.pt
+popd
+
+log "Export to torchscript model"
+./zipformer/export.py \
+  --exp-dir $repo/exp \
+  --use-averaged-model false \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  --epoch 99 \
+  --avg 1 \
+  --jit 1
+
+ls -lh $repo/exp/*.pt
+
+log "Decode with models exported by torch.jit.script()"
+
+./zipformer/jit_pretrained.py \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  --nn-model-filename $repo/exp/jit_script.pt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+for method in greedy_search modified_beam_search fast_beam_search; do
+  log "$method"
+
+  ./zipformer/pretrained.py \
+    --method $method \
+    --beam-size 4 \
+    --checkpoint $repo/exp/pretrained.pt \
+    --tokens $repo/data/lang_bpe_500/tokens.txt \
+    $repo/test_wavs/1089-134686-0001.wav \
+    $repo/test_wavs/1221-135766-0001.wav \
+    $repo/test_wavs/1221-135766-0002.wav
+done
+
+echo "GITHUB_EVENT_NAME: ${GITHUB_EVENT_NAME}"
+echo "GITHUB_EVENT_LABEL_NAME: ${GITHUB_EVENT_LABEL_NAME}"
+if [[ x"${GITHUB_EVENT_NAME}" == x"schedule" || x"${GITHUB_EVENT_LABEL_NAME}" == x"run-decode"  ]]; then
+  mkdir -p zipformer/exp
+  ln -s $PWD/$repo/exp/pretrained.pt zipformer/exp/epoch-999.pt
+  ln -s $PWD/$repo/data/lang_bpe_500 data/
+
+  ls -lh data
+  ls -lh zipformer/exp
+
+  log "Decoding test-clean and test-other"
+
+  # use a small value for decoding with CPU
+  max_duration=100
+
+  for method in greedy_search fast_beam_search modified_beam_search; do
+    log "Decoding with $method"
+
+    ./zipformer/decode.py \
+      --decoding-method $method \
+      --epoch 999 \
+      --avg 1 \
+      --use-averaged-model 0 \
+      --max-duration $max_duration \
+      --exp-dir zipformer/exp
+  done
+
+  rm zipformer/exp/*.pt
+fi
--- a/.github/scripts/run-pre-trained-conformer-ctc.sh
+++ b/.github/scripts/run-pre-trained-conformer-ctc.sh
@ -1,46 +0,0 @@
-#!/usr/bin/env bash
-
-set -e
-
-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]}) $*"
-}
-
-cd egs/librispeech/ASR
-
-repo_url=https://github.com/csukuangfj/icefall-asr-conformer-ctc-bpe-500
-git lfs install
-
-log "Downloading pre-trained model from $repo_url"
-git clone $repo_url
-repo=$(basename $repo_url)
-
-log "Display test files"
-tree $repo/
-ls -lh $repo/test_wavs/*.flac
-
-log "CTC decoding"
-
-./conformer_ctc/pretrained.py \
-  --method ctc-decoding \
-  --num-classes 500 \
-  --checkpoint $repo/exp/pretrained.pt \
-  --tokens $repo/data/lang_bpe_500/tokens.txt \
-  $repo/test_wavs/1089-134686-0001.flac \
-  $repo/test_wavs/1221-135766-0001.flac \
-  $repo/test_wavs/1221-135766-0002.flac
-
-log "HLG decoding"
-
-./conformer_ctc/pretrained.py \
-  --method 1best \
-  --num-classes 500 \
-  --checkpoint $repo/exp/pretrained.pt \
-  --tokens $repo/data/lang_bpe_500/tokens.txt \
-  --words-file $repo/data/lang_bpe_500/words.txt \
-  --HLG $repo/data/lang_bpe_500/HLG.pt \
-  $repo/test_wavs/1089-134686-0001.flac \
-  $repo/test_wavs/1221-135766-0001.flac \
-  $repo/test_wavs/1221-135766-0002.flac
--- a/.github/scripts/run-pre-trained-ctc.sh
+++ b/.github/scripts/run-pre-trained-ctc.sh
@ -0,0 +1,240 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+pushd egs/librispeech/ASR
+
+repo_url=https://huggingface.co/csukuangfj/sherpa-onnx-zipformer-ctc-en-2023-10-02
+log "Downloading pre-trained model from $repo_url"
+git lfs install
+git clone $repo_url
+repo=$(basename $repo_url)
+
+log "Display test files"
+tree $repo/
+ls -lh $repo/test_wavs/*.wav
+
+log "CTC greedy search"
+
+./zipformer/onnx_pretrained_ctc.py \
+  --nn-model $repo/model.onnx \
+  --tokens $repo/tokens.txt \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+log "CTC H decoding"
+
+./zipformer/onnx_pretrained_ctc_H.py \
+  --nn-model $repo/model.onnx \
+  --tokens $repo/tokens.txt \
+  --H $repo/H.fst \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+log "CTC HL decoding"
+
+./zipformer/onnx_pretrained_ctc_HL.py \
+  --nn-model $repo/model.onnx \
+  --words $repo/words.txt \
+  --HL $repo/HL.fst \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+log "CTC HLG decoding"
+
+./zipformer/onnx_pretrained_ctc_HLG.py \
+  --nn-model $repo/model.onnx \
+  --words $repo/words.txt \
+  --HLG $repo/HLG.fst \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+rm -rf $repo
+
+repo_url=https://huggingface.co/csukuangfj/icefall-asr-librispeech-conformer-ctc-jit-bpe-500-2021-11-09
+log "Downloading pre-trained model from $repo_url"
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+pushd $repo
+
+git lfs pull --include "exp/pretrained.pt"
+git lfs pull --include "data/lang_bpe_500/HLG.pt"
+git lfs pull --include "data/lang_bpe_500/L.pt"
+git lfs pull --include "data/lang_bpe_500/L_disambig.pt"
+git lfs pull --include "data/lang_bpe_500/Linv.pt"
+git lfs pull --include "data/lang_bpe_500/bpe.model"
+git lfs pull --include "data/lang_bpe_500/lexicon.txt"
+git lfs pull --include "data/lang_bpe_500/lexicon_disambig.txt"
+git lfs pull --include "data/lang_bpe_500/tokens.txt"
+git lfs pull --include "data/lang_bpe_500/words.txt"
+git lfs pull --include "data/lm/G_3_gram.fst.txt"
+
+popd
+
+log "Display test files"
+tree $repo/
+ls -lh $repo/test_wavs/*.wav
+
+log "CTC decoding"
+
+./conformer_ctc/pretrained.py \
+  --method ctc-decoding \
+  --num-classes 500 \
+  --checkpoint $repo/exp/pretrained.pt \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+log "HLG decoding"
+
+./conformer_ctc/pretrained.py \
+  --method 1best \
+  --num-classes 500 \
+  --checkpoint $repo/exp/pretrained.pt \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  --words-file $repo/data/lang_bpe_500/words.txt \
+  --HLG $repo/data/lang_bpe_500/HLG.pt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+log "CTC decoding on CPU with kaldi decoders using OpenFst"
+
+log "Exporting model with torchscript"
+
+pushd $repo/exp
+ln -s pretrained.pt epoch-99.pt
+popd
+
+./conformer_ctc/export.py \
+  --epoch 99 \
+  --avg 1 \
+  --exp-dir $repo/exp \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  --jit 1
+
+ls -lh $repo/exp
+
+
+log "Generating H.fst, HL.fst"
+
+./local/prepare_lang_fst.py  --lang-dir $repo/data/lang_bpe_500 --ngram-G $repo/data/lm/G_3_gram.fst.txt
+
+ls -lh $repo/data/lang_bpe_500
+
+log "Decoding with H on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_H.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --H $repo/data/lang_bpe_500/H.fst \
+  --tokens $repo/data/lang_bpe_500/tokens.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+log "Decoding with HL on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_HL.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --HL $repo/data/lang_bpe_500/HL.fst \
+  --words $repo/data/lang_bpe_500/words.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+log "Decoding with HLG on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_HLG.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --HLG $repo/data/lang_bpe_500/HLG.fst \
+  --words $repo/data/lang_bpe_500/words.txt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+
+rm -rf $repo
+
+popd
+
+log "Test aishell"
+
+pushd egs/aishell/ASR
+
+repo_url=https://huggingface.co/csukuangfj/icefall_asr_aishell_conformer_ctc
+log "Downloading pre-trained model from $repo_url"
+GIT_LFS_SKIP_SMUDGE=1 git clone $repo_url
+repo=$(basename $repo_url)
+pushd $repo
+
+git lfs pull --include "exp/pretrained.pt"
+git lfs pull --include "data/lang_char/H.fst"
+git lfs pull --include "data/lang_char/HL.fst"
+git lfs pull --include "data/lang_char/HLG.fst"
+
+popd
+
+log "Display test files"
+tree $repo/
+ls -lh $repo/test_wavs/*.wav
+
+log "CTC decoding"
+
+log "Exporting model with torchscript"
+
+pushd $repo/exp
+ln -s pretrained.pt epoch-99.pt
+popd
+
+./conformer_ctc/export.py \
+  --epoch 99 \
+  --avg 1 \
+  --exp-dir $repo/exp \
+  --tokens $repo/data/lang_char/tokens.txt \
+  --jit 1
+
+ls -lh $repo/exp
+
+ls -lh $repo/data/lang_char
+
+log "Decoding with H on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_H.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --H $repo/data/lang_char/H.fst \
+  --tokens $repo/data/lang_char/tokens.txt \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+log "Decoding with HL on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_HL.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --HL $repo/data/lang_char/HL.fst \
+  --words $repo/data/lang_char/words.txt \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+log "Decoding with HLG on CPU with OpenFst"
+
+./conformer_ctc/jit_pretrained_decode_with_HLG.py \
+  --nn-model $repo/exp/cpu_jit.pt \
+  --HLG $repo/data/lang_char/HLG.fst \
+  --words $repo/data/lang_char/words.txt \
+  $repo/test_wavs/0.wav \
+  $repo/test_wavs/1.wav \
+  $repo/test_wavs/2.wav
+
+rm -rf $repo
--- a/.github/scripts/run-swbd-conformer-ctc-2023-08-26.sh
+++ b/.github/scripts/run-swbd-conformer-ctc-2023-08-26.sh
@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+
+set -e
+
+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]}) $*"
+}
+
+cd egs/swbd/ASR
+
+repo_url=https://huggingface.co/zrjin/icefall-asr-swbd-conformer-ctc-2023-8-26
+
+log "Downloading pre-trained model from $repo_url"
+git lfs install
+git clone $repo_url
+repo=$(basename $repo_url)
+
+
+log "Display test files"
+tree $repo/
+ls -lh $repo/test_wavs/*.wav
+
+pushd $repo/exp
+ln -s epoch-98.pt epoch-99.pt
+popd
+
+ls -lh $repo/exp/*.pt
+
+for method in ctc-decoding 1best; do
+  log "$method"
+
+  ./conformer_ctc/pretrained.py \
+    --method $method \
+    --checkpoint $repo/exp/epoch-99.pt \
+    --tokens $repo/data/lang_bpe_500/tokens.txt \
+    --words-file $repo/data/lang_bpe_500/words.txt \
+    --HLG  $repo/data/lang_bpe_500/HLG.pt \
+    --G $repo/data/lm/G_4_gram.pt \
+  $repo/test_wavs/1089-134686-0001.wav \
+  $repo/test_wavs/1221-135766-0001.wav \
+  $repo/test_wavs/1221-135766-0002.wav
+done
--- a/.github/workflows/build-docker-image.yml
+++ b/.github/workflows/build-docker-image.yml
@ -16,7 +16,7 @@ jobs:
      fail-fast: false
      matrix:
        os: [ubuntu-latest]
-        image: ["torch2.0.0-cuda11.7", "torch1.13.0-cuda11.6", "torch1.12.1-cuda11.3", "torch1.9.0-cuda10.2"]
+        image: ["torch2.1.0-cuda12.1", "torch2.1.0-cuda11.8", "torch2.0.0-cuda11.7", "torch1.13.0-cuda11.6", "torch1.12.1-cuda11.3", "torch1.9.0-cuda10.2"]

    steps:
      # refer to https://github.com/actions/checkout
@ -30,6 +30,13 @@ jobs:
          image=${{ matrix.image }}
          mv -v ./docker/$image.dockerfile ./Dockerfile

+      - name: Free space
+        shell: bash
+        run: |
+          df -h
+          rm -rf /opt/hostedtoolcache
+          df -h
+
      - name: Log in to Docker Hub
        uses: docker/login-action@v2
        with:
--- a/.github/workflows/run-docker-image.yml
+++ b/.github/workflows/run-docker-image.yml
@ -14,7 +14,7 @@ jobs:
      fail-fast: false
      matrix:
        os: [ubuntu-latest]
-        image: ["torch2.0.0-cuda11.7", "torch1.13.0-cuda11.6", "torch1.12.1-cuda11.3", "torch1.9.0-cuda10.2"]
+        image: ["torch2.1.0-cuda12.1", "torch2.1.0-cuda11.8", "torch2.0.0-cuda11.7", "torch1.13.0-cuda11.6", "torch1.12.1-cuda11.3", "torch1.9.0-cuda10.2"]
    steps:
      # refer to https://github.com/actions/checkout
      - uses: actions/checkout@v2
@ -30,8 +30,15 @@ jobs:
              uname -a
              cat /etc/*release

+              find / -name libcuda* 2>/dev/null
+
+              ls -lh /usr/local/
+              ls -lh /usr/local/cuda*
+
              nvcc --version

+              ls -lh /usr/local/cuda-*/compat/*
+
              # For torch1.9.0-cuda10.2
              export LD_LIBRARY_PATH=/usr/local/cuda-10.2/compat:$LD_LIBRARY_PATH

@ -41,6 +48,12 @@ jobs:
              # For torch2.0.0-cuda11.7
              export LD_LIBRARY_PATH=/usr/local/cuda-11.7/compat:$LD_LIBRARY_PATH

+              # For torch2.1.0-cuda11.8
+              export LD_LIBRARY_PATH=/usr/local/cuda-11.8/compat:$LD_LIBRARY_PATH
+
+              # For torch2.1.0-cuda12.1
+              export LD_LIBRARY_PATH=/usr/local/cuda-12.1/compat:$LD_LIBRARY_PATH
+

              which nvcc
              cuda_dir=$(dirname $(which nvcc))
--- a/.github/workflows/run-gigaspeech-zipformer-2023-10-17.yml
+++ b/.github/workflows/run-gigaspeech-zipformer-2023-10-17.yml
@ -0,0 +1,126 @@
+# Copyright      2022  Fangjun Kuang (csukuangfj@gmail.com)
+
+# 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.
+
+name: run-gigaspeech-zipformer-2023-10-17
+# zipformer
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    types: [labeled]
+
+  schedule:
+    # minute (0-59)
+    # hour (0-23)
+    # day of the month (1-31)
+    # month (1-12)
+    # day of the week (0-6)
+    # nightly build at 15:50 UTC time every day
+    - cron: "50 15 * * *"
+
+concurrency:
+  group: run_gigaspeech_2023_10_17_zipformer-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  run_gigaspeech_2023_10_17_zipformer:
+    if: github.event.label.name == 'zipformer' ||github.event.label.name == 'ready' || github.event.label.name == 'run-decode' || github.event_name == 'push' || github.event_name == 'schedule'
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: [3.8]
+
+      fail-fast: false
+
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+          cache-dependency-path: '**/requirements-ci.txt'
+
+      - name: Install Python dependencies
+        run: |
+          grep -v '^#' ./requirements-ci.txt  | xargs -n 1 -L 1 pip install
+          pip uninstall -y protobuf
+          pip install --no-binary protobuf protobuf==3.20.*
+
+      - name: Cache kaldifeat
+        id: my-cache
+        uses: actions/cache@v2
+        with:
+          path: |
+            ~/tmp/kaldifeat
+          key: cache-tmp-${{ matrix.python-version }}-2023-05-22
+
+      - name: Install kaldifeat
+        if: steps.my-cache.outputs.cache-hit != 'true'
+        shell: bash
+        run: |
+          .github/scripts/install-kaldifeat.sh
+
+      - name: Inference with pre-trained model
+        shell: bash
+        env:
+          GITHUB_EVENT_NAME: ${{ github.event_name }}
+          GITHUB_EVENT_LABEL_NAME: ${{ github.event.label.name }}
+        run: |
+          mkdir -p egs/gigaspeech/ASR/data
+          ln -sfv ~/tmp/fbank-libri egs/gigaspeech/ASR/data/fbank
+          ls -lh egs/gigaspeech/ASR/data/*
+
+          sudo apt-get -qq install git-lfs tree
+          export PYTHONPATH=$PWD:$PYTHONPATH
+          export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
+          export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
+
+          .github/scripts/run-gigaspeech-zipformer-2023-10-17.sh
+
+      - name: Display decoding results for gigaspeech zipformer
+        if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+        shell: bash
+        run: |
+          cd egs/gigaspeech/ASR/
+          tree ./zipformer/exp
+
+          cd zipformer
+          echo "results for zipformer"
+          echo "===greedy search==="
+          find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+          find exp/greedy_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+          echo "===fast_beam_search==="
+          find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+          find exp/fast_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+          echo "===modified beam search==="
+          find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-clean" {} + | sort -n -k2
+          find exp/modified_beam_search -name "log-*" -exec grep -n --color "best for test-other" {} + | sort -n -k2
+
+      - name: Upload decoding results for gigaspeech zipformer
+        uses: actions/upload-artifact@v2
+        if: github.event_name == 'schedule' || github.event.label.name == 'run-decode'
+        with:
+          name: torch-${{ matrix.torch }}-python-${{ matrix.python-version }}-ubuntu-latest-cpu-zipformer-2022-11-11
+          path: egs/gigaspeech/ASR/zipformer/exp/
--- a/.github/workflows/run-pretrained-conformer-ctc.yml
+++ b/.github/workflows/run-pretrained-conformer-ctc.yml
@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-name: run-pre-trained-conformer-ctc
+name: run-pre-trained-ctc

 on:
  push:
@ -23,13 +23,20 @@ on:
  pull_request:
    types: [labeled]

+  workflow_dispatch:
+    inputs:
+      test-run:
+        description: 'Test (y/n)?'
+        required: true
+        default: 'y'
+
 concurrency:
-  group: run_pre_trained_conformer_ctc-${{ github.ref }}
+  group: run_pre_trained_ctc-${{ github.ref }}
  cancel-in-progress: true

 jobs:
-  run_pre_trained_conformer_ctc:
-    if: github.event.label.name == 'ready' || github.event_name == 'push'
+  run_pre_trained_ctc:
+    if: github.event.label.name == 'ready' || github.event_name == 'push' || github.event.inputs.test-run == 'y' || github.event.label.name == 'ctc'
    runs-on: ${{ matrix.os }}
    strategy:
      matrix:
@ -77,4 +84,4 @@ jobs:
          export PYTHONPATH=$PWD:$PYTHONPATH
          export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
          export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
-          .github/scripts/run-pre-trained-conformer-ctc.sh
+          .github/scripts/run-pre-trained-ctc.sh
--- a/.github/workflows/run-swbd-conformer-ctc.yml
+++ b/.github/workflows/run-swbd-conformer-ctc.yml
@ -0,0 +1,84 @@
+# Copyright      2023   Xiaomi Corp.    (author: Zengrui Jin)
+
+# 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.
+
+name: run-swbd-conformer_ctc
+
+on:
+  push:
+    branches:
+      - master
+  pull_request:
+    types: [labeled]
+
+concurrency:
+  group: run-swbd-conformer_ctc-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  run-swbd-conformer_ctc:
+    if: github.event.label.name == 'onnx' || github.event.label.name == 'ready' || github.event_name == 'push' || github.event.label.name == 'swbd'
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: [3.8]
+
+      fail-fast: false
+
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          fetch-depth: 0
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+          cache-dependency-path: '**/requirements-ci.txt'
+
+      - name: Install Python dependencies
+        run: |
+          grep -v '^#' ./requirements-ci.txt  | xargs -n 1 -L 1 pip install
+          pip uninstall -y protobuf
+          pip install --no-binary protobuf protobuf==3.20.*
+
+      - name: Cache kaldifeat
+        id: my-cache
+        uses: actions/cache@v2
+        with:
+          path: |
+            ~/tmp/kaldifeat
+          key: cache-tmp-${{ matrix.python-version }}-2023-05-22
+
+      - name: Install kaldifeat
+        if: steps.my-cache.outputs.cache-hit != 'true'
+        shell: bash
+        run: |
+          .github/scripts/install-kaldifeat.sh
+
+      - name: Inference with pre-trained model
+        shell: bash
+        env:
+          GITHUB_EVENT_NAME: ${{ github.event_name }}
+          GITHUB_EVENT_LABEL_NAME: ${{ github.event.label.name }}
+        run: |
+          sudo apt-get -qq install git-lfs tree
+          export PYTHONPATH=$PWD:$PYTHONPATH
+          export PYTHONPATH=~/tmp/kaldifeat/kaldifeat/python:$PYTHONPATH
+          export PYTHONPATH=~/tmp/kaldifeat/build/lib:$PYTHONPATH
+
+          .github/scripts/run-swbd-conformer-ctc-2023-08-26.sh
--- a/.github/workflows/run-yesno-recipe.yml
+++ b/.github/workflows/run-yesno-recipe.yml
@ -60,12 +60,12 @@ jobs:

      - name: Install Python dependencies
        run: |
-          grep -v '^#' ./requirements-ci.txt  | grep -v kaldifst | xargs -n 1 -L 1 pip install
+          grep -v '^#' ./requirements-ci.txt  | xargs -n 1 -L 1 pip install
          pip uninstall -y protobuf
          pip install --no-binary protobuf protobuf==3.20.*

-          pip install --no-deps --force-reinstall https://huggingface.co/csukuangfj/k2/resolve/main/cpu/k2-1.24.3.dev20230508+cpu.torch1.13.1-cp38-cp38-linux_x86_64.whl
-          pip install kaldifeat==1.25.0.dev20230726+cpu.torch1.13.1 -f https://csukuangfj.github.io/kaldifeat/cpu.html
+          pip install --no-deps --force-reinstall k2==1.24.4.dev20231021+cpu.torch1.13.1 -f https://k2-fsa.github.io/k2/cpu.html
+          pip install kaldifeat==1.25.1.dev20231022+cpu.torch1.13.1 -f https://csukuangfj.github.io/kaldifeat/cpu.html

      - name: Run yesno recipe
        shell: bash
@ -140,9 +140,46 @@ jobs:
            download/waves_yesno/0_0_0_1_0_0_0_1.wav \
            download/waves_yesno/0_0_1_0_0_0_1_0.wav

+      - name: Test decoding with H
+        shell: bash
+        working-directory: ${{github.workspace}}
+        run: |
+          export PYTHONPATH=$PWD:$PYTHONPATH
+          echo $PYTHONPATH
+
+          cd egs/yesno/ASR
+          python3 ./tdnn/export.py --epoch 14 --avg 2 --jit 1
+
+          python3 ./tdnn/jit_pretrained_decode_with_H.py \
+              --nn-model ./tdnn/exp/cpu_jit.pt \
+              --H ./data/lang_phone/H.fst \
+              --tokens ./data/lang_phone/tokens.txt \
+              ./download/waves_yesno/0_0_0_1_0_0_0_1.wav \
+              ./download/waves_yesno/0_0_1_0_0_0_1_0.wav \
+              ./download/waves_yesno/0_0_1_0_0_1_1_1.wav
+
+      - name: Test decoding with HL
+        shell: bash
+        working-directory: ${{github.workspace}}
+        run: |
+          export PYTHONPATH=$PWD:$PYTHONPATH
+          echo $PYTHONPATH
+
+          cd egs/yesno/ASR
+          python3 ./tdnn/export.py --epoch 14 --avg 2 --jit 1
+
+          python3 ./tdnn/jit_pretrained_decode_with_HL.py \
+              --nn-model ./tdnn/exp/cpu_jit.pt \
+              --HL ./data/lang_phone/HL.fst \
+              --words ./data/lang_phone/words.txt \
+              ./download/waves_yesno/0_0_0_1_0_0_0_1.wav \
+              ./download/waves_yesno/0_0_1_0_0_0_1_0.wav \
+              ./download/waves_yesno/0_0_1_0_0_1_1_1.wav
+
      - name: Show generated files
        shell: bash
        working-directory: ${{github.workspace}}
        run: |
          cd egs/yesno/ASR
          ls -lh tdnn/exp
+          ls -lh data/lang_phone
--- a/.gitignore
+++ b/.gitignore
@ -34,3 +34,5 @@ node_modules
 *.param
 *.bin
 .DS_Store
+*.fst
+*.arpa
--- a/README.md
+++ b/README.md
@ -29,6 +29,7 @@ We provide the following recipes:
  - [yesno][yesno]
  - [LibriSpeech][librispeech]
  - [GigaSpeech][gigaspeech]
+  - [AMI][ami]
  - [Aishell][aishell]
  - [Aishell2][aishell2]
  - [Aishell4][aishell4]
@ -37,6 +38,7 @@ We provide the following recipes:
  - [Aidatatang_200zh][aidatatang_200zh]
  - [WenetSpeech][wenetspeech]
  - [Alimeeting][alimeeting]
+  - [Switchboard][swbd]
  - [TAL_CSASR][tal_csasr]

 ### yesno
@ -118,9 +120,9 @@ We provide a Colab notebook to run a pre-trained transducer conformer + stateles

 | Encoder         | Params | test-clean | test-other |
 |-----------------|--------|------------|------------|
-| zipformer       | 65.5M  | 2.21       | 4.91       |
-| zipformer-small | 23.2M  | 2.46       | 5.83       |
-| zipformer-large | 148.4M | 2.11       | 4.77       |
+| zipformer       | 65.5M  | 2.21       | 4.79       |
+| zipformer-small | 23.2M  | 2.42       | 5.73       |
+| zipformer-large | 148.4M | 2.06       | 4.63       |

 Note: No auxiliary losses are used in the training and no LMs are used
 in the decoding.
@ -146,8 +148,11 @@ in the decoding.

 ### GigaSpeech

-We provide two models for this recipe: [Conformer CTC model][GigaSpeech_conformer_ctc]
-and [Pruned stateless RNN-T: Conformer encoder + Embedding decoder + k2 pruned RNN-T loss][GigaSpeech_pruned_transducer_stateless2].
+We provide three models for this recipe:
+
+- [Conformer CTC model][GigaSpeech_conformer_ctc]
+- [Pruned stateless RNN-T: Conformer encoder + Embedding decoder + k2 pruned RNN-T loss][GigaSpeech_pruned_transducer_stateless2].
+- [Transducer: Zipformer encoder + Embedding decoder][GigaSpeech_zipformer]

 #### Conformer CTC

@ -163,6 +168,14 @@ and [Pruned stateless RNN-T: Conformer encoder + Embedding decoder + k2 pruned R
 |   fast beam search   | 10.50 | 10.69 |
 | modified beam search | 10.40 | 10.51 |

+#### Transducer: Zipformer encoder + Embedding decoder
+
+|                      |  Dev  | Test  |
+|----------------------|-------|-------|
+|    greedy search     | 10.31 | 10.50 |
+|   fast beam search   | 10.26 | 10.48 |
+| modified beam search | 10.25 | 10.38 |
+

 ### Aishell

@ -376,6 +389,7 @@ Please see: [![Open In Colab](https://colab.research.google.com/assets/colab-bad
 [TED-LIUM3_pruned_transducer_stateless]: egs/tedlium3/ASR/pruned_transducer_stateless
 [GigaSpeech_conformer_ctc]: egs/gigaspeech/ASR/conformer_ctc
 [GigaSpeech_pruned_transducer_stateless2]: egs/gigaspeech/ASR/pruned_transducer_stateless2
+[GigaSpeech_zipformer]: egs/gigaspeech/ASR/zipformer
 [Aidatatang_200zh_pruned_transducer_stateless2]: egs/aidatatang_200zh/ASR/pruned_transducer_stateless2
 [WenetSpeech_pruned_transducer_stateless2]: egs/wenetspeech/ASR/pruned_transducer_stateless2
 [WenetSpeech_pruned_transducer_stateless5]: egs/wenetspeech/ASR/pruned_transducer_stateless5
@ -393,4 +407,6 @@ Please see: [![Open In Colab](https://colab.research.google.com/assets/colab-bad
 [wenetspeech]: egs/wenetspeech/ASR
 [alimeeting]: egs/alimeeting/ASR
 [tal_csasr]: egs/tal_csasr/ASR
+[ami]: egs/ami
+[swbd]: egs/swbd/ASR
 [k2]: https://github.com/k2-fsa/k2
--- a/docker/torch1.12.1-cuda11.3.dockerfile
+++ b/docker/torch1.12.1-cuda11.3.dockerfile
@ -4,8 +4,9 @@ ENV LC_ALL C.UTF-8

 ARG DEBIAN_FRONTEND=noninteractive

-ARG K2_VERSION="1.24.3.dev20230725+cuda11.3.torch1.12.1"
-ARG KALDIFEAT_VERSION="1.25.0.dev20230726+cuda11.3.torch1.12.1"
+# python 3.7
+ARG K2_VERSION="1.24.4.dev20230725+cuda11.3.torch1.12.1"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda11.3.torch1.12.1"
 ARG TORCHAUDIO_VERSION="0.12.1+cu113"

 LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
--- a/docker/torch1.13.0-cuda11.6.dockerfile
+++ b/docker/torch1.13.0-cuda11.6.dockerfile
@ -4,8 +4,9 @@ ENV LC_ALL C.UTF-8

 ARG DEBIAN_FRONTEND=noninteractive

-ARG K2_VERSION="1.24.3.dev20230725+cuda11.6.torch1.13.0"
-ARG KALDIFEAT_VERSION="1.25.0.dev20230726+cuda11.6.torch1.13.0"
+# python 3.9
+ARG K2_VERSION="1.24.4.dev20231021+cuda11.6.torch1.13.0"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda11.6.torch1.13.0"
 ARG TORCHAUDIO_VERSION="0.13.0+cu116"

 LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
--- a/docker/torch1.9.0-cuda10.2.dockerfile
+++ b/docker/torch1.9.0-cuda10.2.dockerfile
@ -4,8 +4,9 @@ ENV LC_ALL C.UTF-8

 ARG DEBIAN_FRONTEND=noninteractive

+# python 3.7
 ARG K2_VERSION="1.24.3.dev20230726+cuda10.2.torch1.9.0"
-ARG KALDIFEAT_VERSION="1.25.0.dev20230726+cuda10.2.torch1.9.0"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda10.2.torch1.9.0"
 ARG TORCHAUDIO_VERSION="0.9.0"

 LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
--- a/docker/torch2.0.0-cuda11.7.dockerfile
+++ b/docker/torch2.0.0-cuda11.7.dockerfile
@ -4,8 +4,9 @@ ENV LC_ALL C.UTF-8

 ARG DEBIAN_FRONTEND=noninteractive

-ARG K2_VERSION="1.24.3.dev20230718+cuda11.7.torch2.0.0"
-ARG KALDIFEAT_VERSION="1.25.0.dev20230726+cuda11.7.torch2.0.0"
+# python 3.10
+ARG K2_VERSION="1.24.4.dev20231021+cuda11.7.torch2.0.0"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda11.7.torch2.0.0"
 ARG TORCHAUDIO_VERSION="2.0.0+cu117"

 LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
--- a/docker/torch2.1.0-cuda11.8.dockerfile
+++ b/docker/torch2.1.0-cuda11.8.dockerfile
@ -0,0 +1,71 @@
+FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-devel
+
+ENV LC_ALL C.UTF-8
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+# python 3.10
+ARG K2_VERSION="1.24.4.dev20231021+cuda11.8.torch2.1.0"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda11.8.torch2.1.0"
+ARG TORCHAUDIO_VERSION="2.1.0+cu118"
+
+LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
+LABEL k2_version=${K2_VERSION}
+LABEL kaldifeat_version=${KALDIFEAT_VERSION}
+LABEL github_repo="https://github.com/k2-fsa/icefall"
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        curl \
+        vim \
+    		libssl-dev \
+        autoconf \
+        automake \
+        bzip2 \
+        ca-certificates \
+        ffmpeg \
+        g++ \
+        gfortran \
+        git \
+        libtool \
+        make \
+        patch \
+        sox \
+        subversion \
+        unzip \
+        valgrind \
+        wget \
+        zlib1g-dev \
+        && rm -rf /var/lib/apt/lists/*
+
+# Install dependencies
+RUN pip install --no-cache-dir \
+      torchaudio==${TORCHAUDIO_VERSION} -f https://download.pytorch.org/whl/torch_stable.html \
+      k2==${K2_VERSION} -f https://k2-fsa.github.io/k2/cuda.html \
+      git+https://github.com/lhotse-speech/lhotse \
+      kaldifeat==${KALDIFEAT_VERSION} -f https://csukuangfj.github.io/kaldifeat/cuda.html \
+      \
+      kaldi_native_io \
+      kaldialign \
+      kaldifst \
+      kaldilm \
+      sentencepiece>=0.1.96 \
+      tensorboard \
+      typeguard \
+      dill \
+      onnx \
+      onnxruntime \
+      onnxmltools \
+      multi_quantization \
+      typeguard \
+      numpy \
+      pytest \
+      graphviz
+
+RUN git clone https://github.com/k2-fsa/icefall /workspace/icefall && \
+    cd /workspace/icefall && \
+    pip install --no-cache-dir -r requirements.txt
+
+ENV PYTHONPATH /workspace/icefall:$PYTHONPATH
+
+WORKDIR /workspace/icefall
--- a/docker/torch2.1.0-cuda12.1.dockerfile
+++ b/docker/torch2.1.0-cuda12.1.dockerfile
@ -0,0 +1,71 @@
+FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel
+
+ENV LC_ALL C.UTF-8
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+# python 3.10
+ARG K2_VERSION="1.24.4.dev20231021+cuda12.1.torch2.1.0"
+ARG KALDIFEAT_VERSION="1.25.1.dev20231022+cuda12.1.torch2.1.0"
+ARG TORCHAUDIO_VERSION="2.1.0+cu121"
+
+LABEL authors="Fangjun Kuang <csukuangfj@gmail.com>"
+LABEL k2_version=${K2_VERSION}
+LABEL kaldifeat_version=${KALDIFEAT_VERSION}
+LABEL github_repo="https://github.com/k2-fsa/icefall"
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+        curl \
+        vim \
+    		libssl-dev \
+        autoconf \
+        automake \
+        bzip2 \
+        ca-certificates \
+        ffmpeg \
+        g++ \
+        gfortran \
+        git \
+        libtool \
+        make \
+        patch \
+        sox \
+        subversion \
+        unzip \
+        valgrind \
+        wget \
+        zlib1g-dev \
+        && rm -rf /var/lib/apt/lists/*
+
+# Install dependencies
+RUN pip install --no-cache-dir \
+      torchaudio==${TORCHAUDIO_VERSION} -f https://download.pytorch.org/whl/torch_stable.html \
+      k2==${K2_VERSION} -f https://k2-fsa.github.io/k2/cuda.html \
+      git+https://github.com/lhotse-speech/lhotse \
+      kaldifeat==${KALDIFEAT_VERSION} -f https://csukuangfj.github.io/kaldifeat/cuda.html \
+      \
+      kaldi_native_io \
+      kaldialign \
+      kaldifst \
+      kaldilm \
+      sentencepiece>=0.1.96 \
+      tensorboard \
+      typeguard \
+      dill \
+      onnx \
+      onnxruntime \
+      onnxmltools \
+      multi_quantization \
+      typeguard \
+      numpy \
+      pytest \
+      graphviz
+
+RUN git clone https://github.com/k2-fsa/icefall /workspace/icefall && \
+    cd /workspace/icefall && \
+    pip install --no-cache-dir -r requirements.txt
+
+ENV PYTHONPATH /workspace/icefall:$PYTHONPATH
+
+WORKDIR /workspace/icefall
--- a/docs/source/contributing/code-style.rst
+++ b/docs/source/contributing/code-style.rst
@ -38,7 +38,7 @@ Please fix any issues reported by the check tools.
 .. HINT::

  Some of the check tools, i.e., ``black`` and ``isort`` will modify
-  the files to be commited **in-place**. So please run ``git status``
+  the files to be committed **in-place**. So please run ``git status``
  after failure to see which file has been modified by the tools
  before you make any further changes.

--- a/docs/source/decoding-with-langugage-models/LODR.rst
+++ b/docs/source/decoding-with-langugage-models/LODR.rst
@ -56,7 +56,7 @@ during decoding for transducer model:
    \lambda_1 \log p_{\text{Target LM}}\left(y_u|\mathit{x},y_{1:u-1}\right) -
    \lambda_2 \log p_{\text{bi-gram}}\left(y_u|\mathit{x},y_{1:u-1}\right)

-In LODR, an additional bi-gram LM estimated on the source domain (e.g training corpus) is required. Comared to DR,
+In LODR, an additional bi-gram LM estimated on the source domain (e.g training corpus) is required. Compared to DR,
 the only difference lies in the choice of source domain LM. According to the original `paper <https://arxiv.org/abs/2203.16776>`_,
 LODR achieves similar performance compared DR in both intra-domain and cross-domain settings.
 As a bi-gram is much faster to evaluate, LODR is usually much faster.
--- a/docs/source/decoding-with-langugage-models/index.rst
+++ b/docs/source/decoding-with-langugage-models/index.rst
@ -2,12 +2,13 @@ Decoding with language models
 =============================

 This section describes how to use external langugage models 
-during decoding to improve the WER of transducer models.
+during decoding to improve the WER of transducer models. To train an external language model,
+please refer to this tutorial: :ref:`train_nnlm`.

 The following decoding methods with external langugage models are available:


-.. list-table:: LM-rescoring-based methods vs shallow-fusion-based methods (The numbers in each field is WER on test-clean, WER on test-other and decoding time on test-clean)
+.. list-table:: 
   :widths: 25 50
   :header-rows: 1

--- a/docs/source/docker/intro.rst
+++ b/docs/source/docker/intro.rst
@ -30,6 +30,8 @@ which will give you something like below:

 .. code-block:: bash

+  "torch2.1.0-cuda12.1"
+  "torch2.1.0-cuda11.8"
  "torch2.0.0-cuda11.7"
  "torch1.12.1-cuda11.3"
  "torch1.9.0-cuda10.2"
--- a/docs/source/model-export/export-ncnn-conv-emformer.rst
+++ b/docs/source/model-export/export-ncnn-conv-emformer.rst
@ -125,7 +125,7 @@ Python code. We have also set up ``PATH`` so that you can use
 .. caution::

  Please don't use `<https://github.com/tencent/ncnn>`_.
-  We have made some modifications to the offical `ncnn`_.
+  We have made some modifications to the official `ncnn`_.

  We will synchronize `<https://github.com/csukuangfj/ncnn>`_ periodically
  with the official one.
--- a/docs/source/model-export/export-ncnn.rst
+++ b/docs/source/model-export/export-ncnn.rst
@ -1,3 +1,5 @@
+.. _icefall_export_to_ncnn:
+
 Export to ncnn
 ==============

--- a/docs/source/recipes/Non-streaming-ASR/librispeech/distillation.rst
+++ b/docs/source/recipes/Non-streaming-ASR/librispeech/distillation.rst
@ -47,7 +47,7 @@ The data preparation contains several stages, you can use the following two
 options:

  - ``--stage``
-  - ``--stop-stage``
+  - ``--stop_stage``

 to control which stage(s) should be run. By default, all stages are executed.

@ -56,8 +56,8 @@ For example,
 .. code-block:: bash

  $ cd egs/librispeech/ASR
-  $ ./prepare.sh --stage 0 --stop-stage 0 # run only stage 0
-  $ ./prepare.sh --stage 2 --stop-stage 5 # run from stage 2 to stage 5
+  $ ./prepare.sh --stage 0 --stop_stage 0 # run only stage 0
+  $ ./prepare.sh --stage 2 --stop_stage 5 # run from stage 2 to stage 5

 .. HINT::

@ -108,15 +108,15 @@ As usual, you can control the stages you want to run by specifying the following
 two options:

  - ``--stage``
-  - ``--stop-stage``
+  - ``--stop_stage``

 For example,

 .. code-block:: bash

  $ cd egs/librispeech/ASR
-  $ ./distillation_with_hubert.sh --stage 0 --stop-stage 0 # run only stage 0
-  $ ./distillation_with_hubert.sh --stage 2 --stop-stage 4 # run from stage 2 to stage 5
+  $ ./distillation_with_hubert.sh --stage 0 --stop_stage 0 # run only stage 0
+  $ ./distillation_with_hubert.sh --stage 2 --stop_stage 4 # run from stage 2 to stage 5

 Here are a few options in `./distillation_with_hubert.sh <https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/distillation_with_hubert.sh>`_
 you need to know before you proceed.
@ -134,7 +134,7 @@ and prepares MVQ-augmented training manifests.

 .. code-block:: bash

-  $ ./distillation_with_hubert.sh --stage 2 --stop-stage 2 # run only stage 2
+  $ ./distillation_with_hubert.sh --stage 2 --stop_stage 2 # run only stage 2

 Please see the
 following screenshot for the output of an example execution.
@ -172,7 +172,7 @@ To perform training, please run stage 3 by executing the following command.

 .. code-block:: bash

-  $ ./prepare.sh --stage 3 --stop-stage 3 # run MVQ training
+  $ ./prepare.sh --stage 3 --stop_stage 3 # run MVQ training

 Here is the code snippet for training:

--- a/docs/source/recipes/RNN-LM/index.rst
+++ b/docs/source/recipes/RNN-LM/index.rst
@ -0,0 +1,7 @@
+RNN-LM
+======
+
+.. toctree::
+   :maxdepth: 2
+
+   librispeech/lm-training
--- a/docs/source/recipes/RNN-LM/librispeech/lm-training.rst
+++ b/docs/source/recipes/RNN-LM/librispeech/lm-training.rst
@ -0,0 +1,104 @@
+.. _train_nnlm:
+
+Train an RNN langugage model
+======================================
+
+If you have enough text data, you can train a neural network language model (NNLM) to improve
+the WER of your E2E ASR system. This tutorial shows you how to train an RNNLM from 
+scratch.
+
+.. HINT::
+
+    For how to use an NNLM during decoding, please refer to the following tutorials:
+    :ref:`shallow_fusion`, :ref:`LODR`, :ref:`rescoring`
+
+.. note::
+
+    This tutorial is based on the LibriSpeech recipe. Please check it out for the necessary
+    python scripts for this tutorial. We use the LibriSpeech LM-corpus as the LM training set 
+    for illustration purpose. You can also collect your own data. The data format is quite simple:
+    each line should contain a complete sentence, and words should be separated by space.
+
+First, let's download the training data for the RNNLM. This can be done via the 
+following command:
+
+.. code-block:: bash
+
+    $ wget https://www.openslr.org/resources/11/librispeech-lm-norm.txt.gz 
+    $ gzip -d librispeech-lm-norm.txt.gz
+
+As we are training a BPE-level RNNLM, we need to tokenize the training text, which requires a
+BPE tokenizer. This can be achieved by executing the following command:
+
+.. code-block:: bash
+    
+    $ # if you don't have the BPE
+    $ GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/Zengwei/icefall-asr-librispeech-zipformer-2023-05-15
+    $ cd icefall-asr-librispeech-zipformer-2023-05-15/data/lang_bpe_500
+    $ git lfs pull --include bpe.model
+    $ cd ../../..
+
+    $ ./local/prepare_lm_training_data.py \
+        --bpe-model icefall-asr-librispeech-zipformer-2023-05-15/data/lang_bpe_500/bpe.model \
+        --lm-data librispeech-lm-norm.txt \
+        --lm-archive data/lang_bpe_500/lm_data.pt
+
+Now, you should have a file name ``lm_data.pt`` file store under the directory ``data/lang_bpe_500``.
+This is the packed training data for the RNNLM. We then sort the training data according to its
+sentence length.
+
+.. code-block:: bash
+
+    $ # This could take a while (~ 20 minutes), feel free to grab a cup of coffee :)
+    $ ./local/sort_lm_training_data.py \
+        --in-lm-data data/lang_bpe_500/lm_data.pt \
+        --out-lm-data data/lang_bpe_500/sorted_lm_data.pt \
+        --out-statistics data/lang_bpe_500/lm_data_stats.txt
+
+
+The aforementioned steps can be repeated to create a a validation set for you RNNLM. Let's say 
+you have a validation set in ``valid.txt``, you can just set ``--lm-data valid.txt`` 
+and ``--lm-archive data/lang_bpe_500/lm-data-valid.pt`` when calling ``./local/prepare_lm_training_data.py``.
+
+After completing the previous steps, the training and testing sets for training RNNLM are ready. 
+The next step is to train the RNNLM model. The training command is as follows:
+
+.. code-block:: bash
+
+    $ # assume you are in the icefall root directory
+    $ cd rnn_lm
+    $ ln -s ../../egs/librispeech/ASR/data .
+    $ cd ..
+    $ ./rnn_lm/train.py \
+        --world-size 4 \
+        --exp-dir ./rnn_lm/exp \
+        --start-epoch 0 \
+        --num-epochs 10 \
+        --use-fp16 0 \
+        --tie-weights 1 \
+        --embedding-dim 2048 \
+        --hidden_dim 2048 \
+        --num-layers 3 \
+        --batch-size 300 \
+        --lm-data rnn_lm/data/lang_bpe_500/sorted_lm_data.pt \
+        --lm-data-valid rnn_lm/data/lang_bpe_500/sorted_lm_data.pt
+
+
+.. note::
+
+    You can adjust the RNNLM hyper parameters to control the size of the RNNLM,
+    such as embedding dimension and hidden state dimension. For more details, please
+    run ``./rnn_lm/train.py --help``.
+
+.. note::
+
+    The training of RNNLM can take a long time (usually a couple of days).
+
+
+
+
+
+
+
+
+
--- a/docs/source/recipes/index.rst
+++ b/docs/source/recipes/index.rst
@ -15,3 +15,4 @@ We may add recipes for other tasks as well in the future.

   Non-streaming-ASR/index
   Streaming-ASR/index
+   RNN-LM/index
--- a/egs/aidatatang_200zh/ASR/pruned_transducer_stateless2/asr_datamodule.py
+++ b/egs/aidatatang_200zh/ASR/pruned_transducer_stateless2/asr_datamodule.py
@ -211,7 +211,7 @@ class Aidatatang_200zhAsrDataModule:
        if self.args.enable_musan:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/aidatatang_200zh/ASR/pruned_transducer_stateless2/train.py
+++ b/egs/aidatatang_200zh/ASR/pruned_transducer_stateless2/train.py
@ -635,7 +635,6 @@ def train_one_epoch(
    tot_loss = MetricsTracker()

    for batch_idx, batch in enumerate(train_dl):
-
        params.batch_idx_train += 1
        batch_size = len(batch["supervisions"]["text"])

@ -800,7 +799,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/RESULTS.md
+++ b/egs/aishell/ASR/RESULTS.md
@ -145,6 +145,56 @@ for m in greedy_search modified_beam_search fast_beam_search ; do
 done
 ```

+#### Pruned transducer stateless 7 streaming
+[./pruned_transducer_stateless7_streaming](./pruned_transducer_stateless7_streaming)
+
+It's Streaming version of Zipformer1 with Pruned RNNT loss.
+
+|                        | test | dev  | comment                               |
+|------------------------|------|------|---------------------------------------|
+| greedy search          | 6.95 | 6.29 | --epoch 44 --avg 15 --max-duration 600 |
+| modified beam search   | 6.51 | 5.90 | --epoch 44 --avg 15 --max-duration 600 |
+| fast beam search       | 6.73 | 6.09 | --epoch 44 --avg 15 --max-duration 600 |
+
+Training command is:
+
+```bash
+./prepare.sh
+
+export CUDA_VISIBLE_DEVICES="0,1"
+
+./pruned_transducer_stateless7_streaming/train.py \
+    --world-size 2 \
+    --num-epochs 50 \
+    --use-fp16 1 \
+    --context-size 1 \
+    --max-duration 800 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --enable-musan 0 \
+    --spec-aug-time-warp-factor 20
+```
+
+**Caution**: It uses `--context-size=1`.
+
+The decoding command is:
+```bash
+for m in greedy_search modified_beam_search fast_beam_search ; do
+  ./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 44 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --lang-dir data/lang_char \
+    --context-size 1 \
+    --decoding-method $m
+done
+```
+
+Pretrained models, training logs, decoding logs, tensorboard and decoding results
+are available at
+<https://huggingface.co/zrjin/icefall-asr-aishell-zipformer-pruned-transducer-stateless7-streaming-2023-10-16/>
+
+
+
 #### Pruned transducer stateless 7

 [./pruned_transducer_stateless7](./pruned_transducer_stateless7)
--- a/egs/aishell/ASR/conformer_ctc/export.py
+++ b/egs/aishell/ASR/conformer_ctc/export.py
@ -23,12 +23,12 @@ import argparse
 import logging
 from pathlib import Path

+import k2
 import torch
 from conformer import Conformer

 from icefall.checkpoint import average_checkpoints, load_checkpoint
-from icefall.lexicon import Lexicon
-from icefall.utils import AttributeDict, str2bool
+from icefall.utils import AttributeDict, num_tokens, str2bool


 def get_parser():
@ -63,11 +63,10 @@ def get_parser():
    )

    parser.add_argument(
-        "--lang-dir",
+        "--tokens",
        type=str,
-        default="data/lang_char",
-        help="""It contains language related input files such as "lexicon.txt"
-        """,
+        required=True,
+        help="Path to the tokens.txt.",
    )

    parser.add_argument(
@ -98,16 +97,16 @@ def get_params() -> AttributeDict:
 def main():
    args = get_parser().parse_args()
    args.exp_dir = Path(args.exp_dir)
-    args.lang_dir = Path(args.lang_dir)

    params = get_params()
    params.update(vars(args))

-    logging.info(params)
+    # Load tokens.txt here
+    token_table = k2.SymbolTable.from_file(params.tokens)

-    lexicon = Lexicon(params.lang_dir)
-    max_token_id = max(lexicon.tokens)
-    num_classes = max_token_id + 1  # +1 for the blank
+    num_classes = num_tokens(token_table) + 1  # +1 for the blank
+
+    logging.info(params)

    device = torch.device("cpu")
    if torch.cuda.is_available():
--- a/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_H.py
+++ b/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_H.py
@ -0,0 +1 @@
+../../../librispeech/ASR/conformer_ctc/jit_pretrained_decode_with_H.py
--- a/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_HL.py
+++ b/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_HL.py
@ -0,0 +1 @@
+../../../librispeech/ASR/conformer_ctc/jit_pretrained_decode_with_HL.py
--- a/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_HLG.py
+++ b/egs/aishell/ASR/conformer_ctc/jit_pretrained_decode_with_HLG.py
@ -0,0 +1 @@
+../../../librispeech/ASR/conformer_ctc/jit_pretrained_decode_with_HLG.py
--- a/egs/aishell/ASR/conformer_ctc/test_transformer.py
+++ b/egs/aishell/ASR/conformer_ctc/test_transformer.py
--- a/egs/aishell/ASR/local/prepare_lang_fst.py
+++ b/egs/aishell/ASR/local/prepare_lang_fst.py
@ -0,0 +1 @@
+../../../librispeech/ASR/local/prepare_lang_fst.py
--- a/egs/aishell/ASR/local/train_bbpe_model.py
+++ b/egs/aishell/ASR/local/train_bbpe_model.py
@ -15,7 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

-
 # You can install sentencepiece via:
 #
 #  pip install sentencepiece
@ -26,12 +25,12 @@
 # Please install a version >=0.1.96

 import argparse
-import re
 import shutil
 import tempfile
 from pathlib import Path

 import sentencepiece as spm
+
 from icefall import byte_encode, tokenize_by_CJK_char


@ -74,6 +73,11 @@ def main():
    model_type = "unigram"

    model_prefix = f"{lang_dir}/{model_type}_{vocab_size}"
+    model_file = Path(model_prefix + ".model")
+    if model_file.is_file():
+        print(f"{model_file} exists - skipping")
+        return
+
    character_coverage = 1.0
    input_sentence_size = 100000000

@ -88,23 +92,18 @@ def main():

    _convert_to_bchar(args.transcript, train_text)

-    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,
-            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,
-        )
-    else:
-        print(f"{model_file} exists - skipping")
-        return
+    spm.SentencePieceTrainer.train(
+        input=train_text,
+        vocab_size=vocab_size,
+        model_type=model_type,
+        model_prefix=model_prefix,
+        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,
+    )

    shutil.copyfile(model_file, f"{lang_dir}/bbpe.model")

--- a/egs/aishell/ASR/prepare.sh
+++ b/egs/aishell/ASR/prepare.sh
@ -144,6 +144,7 @@ if [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then
    ./local/prepare_lang.py --lang-dir $lang_phone_dir
  fi

+
  # Train a bigram P for MMI training
  if [ ! -f $lang_phone_dir/transcript_words.txt ]; then
    log "Generate data to train phone based bigram P"
@ -258,6 +259,13 @@ if [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then
      --max-order=3 \
      data/lm/3-gram.unpruned.arpa > data/lm/G_3_gram_char.fst.txt
  fi
+
+  if [ ! -f $lang_char_dir/HLG.fst ]; then
+    lang_phone_dir=data/lang_phone
+    ./local/prepare_lang_fst.py  \
+      --lang-dir $lang_phone_dir \
+      --ngram-G ./data/lm/G_3_gram.fst.txt
+  fi
 fi

 if [ $stage -le 9 ] && [ $stop_stage -ge 9 ]; then
--- a/egs/aishell/ASR/pruned_transducer_stateless2/train.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless2/train.py
@ -872,7 +872,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/pruned_transducer_stateless3/train.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless3/train.py
@ -1045,7 +1045,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/pruned_transducer_stateless7/train.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7/train.py
@ -703,7 +703,7 @@ def compute_loss(
            if batch_idx_train >= warm_step
            else 0.1 + 0.9 * (batch_idx_train / warm_step)
        )
-        loss = params.simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss
+        loss = simple_loss_scale * simple_loss + pruned_loss_scale * pruned_loss

    assert loss.requires_grad == is_training

@ -1028,7 +1028,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/pruned_transducer_stateless7/train2.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7/train2.py
@ -1031,7 +1031,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/pruned_transducer_stateless7_bbpe/train.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_bbpe/train.py
@ -1019,7 +1019,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/README.md
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/README.md
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/README.md
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/asr_datamodule.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/asr_datamodule.py
@ -0,0 +1 @@
+../pruned_transducer_stateless2/asr_datamodule.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/beam_search.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/beam_search.py
@ -0,0 +1 @@
+../pruned_transducer_stateless2/beam_search.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decode.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decode.py
@ -0,0 +1,735 @@
+#!/usr/bin/env python3
+#
+# Copyright 2021-2022 Xiaomi Corporation (Author: Fangjun Kuang,
+#                                                 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.
+"""
+Usage:
+(1) greedy search
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method greedy_search
+
+(2) beam search (not recommended)
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method beam_search \
+    --beam-size 4
+
+(3) modified beam search
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method modified_beam_search \
+    --beam-size 4
+
+(4) fast beam search (one best)
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method fast_beam_search \
+    --beam 20.0 \
+    --max-contexts 8 \
+    --max-states 64
+
+(5) fast beam search (nbest)
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method fast_beam_search_nbest \
+    --beam 20.0 \
+    --max-contexts 8 \
+    --max-states 64 \
+    --num-paths 200 \
+    --nbest-scale 0.5
+
+(6) fast beam search (nbest oracle WER)
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method fast_beam_search_nbest_oracle \
+    --beam 20.0 \
+    --max-contexts 8 \
+    --max-states 64 \
+    --num-paths 200 \
+    --nbest-scale 0.5
+
+(7) fast beam search (with LG)
+./pruned_transducer_stateless7_streaming/decode.py \
+    --epoch 28 \
+    --avg 15 \
+    --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+    --max-duration 600 \
+    --decode-chunk-len 32 \
+    --decoding-method fast_beam_search_nbest_LG \
+    --beam 20.0 \
+    --max-contexts 8 \
+    --max-states 64
+"""
+
+
+import argparse
+import logging
+import math
+import os
+from collections import defaultdict
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import torch
+import torch.nn as nn
+from asr_datamodule import AishellAsrDataModule
+from beam_search import (
+    beam_search,
+    fast_beam_search_one_best,
+    greedy_search,
+    greedy_search_batch,
+    modified_beam_search,
+)
+from train import add_model_arguments, get_params, get_transducer_model
+
+from icefall import ContextGraph
+from icefall.char_graph_compiler import CharCtcTrainingGraphCompiler
+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 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=False,
+        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(
+        "--exp-dir",
+        type=str,
+        default="pruned_transducer_stateless3/exp",
+        help="The experiment dir",
+    )
+
+    parser.add_argument(
+        "--lang-dir",
+        type=str,
+        default="data/lang_char",
+        help="The lang dir",
+    )
+
+    parser.add_argument(
+        "--decoding-method",
+        type=str,
+        default="greedy_search",
+        help="""Possible values are:
+          - greedy_search
+          - beam_search
+          - modified_beam_search
+          - fast_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=4,
+        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""",
+    )
+
+    parser.add_argument(
+        "--max-contexts",
+        type=int,
+        default=4,
+        help="""Used only when --decoding-method is
+        fast_beam_search""",
+    )
+
+    parser.add_argument(
+        "--max-states",
+        type=int,
+        default=8,
+        help="""Used only when --decoding-method is
+        fast_beam_search""",
+    )
+
+    parser.add_argument(
+        "--context-size",
+        type=int,
+        default=1,
+        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(
+        "--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.
+        """,
+    )
+
+    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.
+        """,
+    )
+
+    add_model_arguments(parser)
+
+    return parser
+
+
+def decode_one_batch(
+    params: AttributeDict,
+    model: nn.Module,
+    token_table: k2.SymbolTable,
+    batch: dict,
+    decoding_graph: Optional[k2.Fsa] = None,
+    context_graph: Optional[ContextGraph] = 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.
+      token_table:
+        It maps token ID to a string.
+      batch:
+        It is the return value from iterating
+        `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation
+        for the format of the `batch`.
+      decoding_graph:
+        The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+        only when --decoding_method is fast_beam_search.
+    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)
+
+    feature_lens += 30
+    feature = torch.nn.functional.pad(
+        feature,
+        pad=(0, 0, 0, 30),
+        value=LOG_EPS,
+    )
+    encoder_out, encoder_out_lens = model.encoder(x=feature, x_lens=feature_lens)
+
+    if params.decoding_method == "fast_beam_search":
+        hyp_tokens = fast_beam_search_one_best(
+            model=model,
+            decoding_graph=decoding_graph,
+            encoder_out=encoder_out,
+            encoder_out_lens=encoder_out_lens,
+            beam=params.beam,
+            max_contexts=params.max_contexts,
+            max_states=params.max_states,
+        )
+    elif params.decoding_method == "greedy_search" and params.max_sym_per_frame == 1:
+        hyp_tokens = greedy_search_batch(
+            model=model,
+            encoder_out=encoder_out,
+            encoder_out_lens=encoder_out_lens,
+        )
+    elif params.decoding_method == "modified_beam_search":
+        hyp_tokens = modified_beam_search(
+            model=model,
+            encoder_out=encoder_out,
+            encoder_out_lens=encoder_out_lens,
+            beam=params.beam_size,
+            context_graph=context_graph,
+        )
+    else:
+        hyp_tokens = []
+        batch_size = encoder_out.size(0)
+        for i in range(batch_size):
+            # fmt: off
+            encoder_out_i = encoder_out[i:i+1, :encoder_out_lens[i]]
+            # fmt: on
+            if params.decoding_method == "greedy_search":
+                hyp = greedy_search(
+                    model=model,
+                    encoder_out=encoder_out_i,
+                    max_sym_per_frame=params.max_sym_per_frame,
+                )
+            elif params.decoding_method == "beam_search":
+                hyp = beam_search(
+                    model=model,
+                    encoder_out=encoder_out_i,
+                    beam=params.beam_size,
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported decoding method: {params.decoding_method}"
+                )
+            hyp_tokens.append(hyp)
+
+    hyps = [[token_table[t] for t in tokens] for tokens in hyp_tokens]
+
+    if params.decoding_method == "greedy_search":
+        return {"greedy_search": hyps}
+    elif params.decoding_method == "fast_beam_search":
+        return {
+            (
+                f"beam_{params.beam}_"
+                f"max_contexts_{params.max_contexts}_"
+                f"max_states_{params.max_states}"
+            ): hyps
+        }
+    else:
+        key = f"beam_size_{params.beam_size}"
+        if params.has_contexts:
+            key += f"-context-score-{params.context_score}"
+        else:
+            key += "-no-context-words"
+        return {key: hyps}
+
+
+def decode_dataset(
+    dl: torch.utils.data.DataLoader,
+    params: AttributeDict,
+    model: nn.Module,
+    token_table: Optional[k2.SymbolTable] = None,
+    decoding_graph: Optional[k2.Fsa] = None,
+    context_graph: Optional[ContextGraph] = None,
+) -> 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.
+      token_table:
+        It maps a token ID to a string.
+      decoding_graph:
+        The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+        only when --decoding_method is fast_beam_search.
+    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 = defaultdict(list)
+    for batch_idx, batch in enumerate(dl):
+        texts = batch["supervisions"]["text"]
+        cut_ids = [cut.id for cut in batch["supervisions"]["cut"]]
+
+        hyps_dict = decode_one_batch(
+            params=params,
+            model=model,
+            token_table=token_table,
+            decoding_graph=decoding_graph,
+            context_graph=context_graph,
+            batch=batch,
+        )
+
+        for name, hyps in hyps_dict.items():
+            this_batch = []
+            assert len(hyps) == len(texts)
+            for cut_id, hyp_words, ref_text in zip(cut_ids, hyps, texts):
+                ref_words = ref_text.split()
+                this_batch.append((cut_id, ref_words, hyp_words))
+
+            results[name].extend(this_batch)
+
+        num_cuts += len(texts)
+
+        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}")
+    return results
+
+
+def save_results(
+    params: AttributeDict,
+    test_set_name: str,
+    results_dict: Dict[str, List[Tuple[str, List[str], List[str]]]],
+):
+    test_set_wers = dict()
+    for key, results in results_dict.items():
+        recog_path = params.res_dir / f"recogs-{test_set_name}-{params.suffix}.txt"
+        results = sorted(results)
+        store_transcripts(filename=recog_path, texts=results)
+        logging.info(f"The transcripts are stored in {recog_path}")
+
+        # The following prints out WERs, per-word error statistics and aligned
+        # ref/hyp pairs.
+        errs_filename = params.res_dir / f"errs-{test_set_name}-{params.suffix}.txt"
+        with open(errs_filename, "w") as f:
+            wer = write_error_stats(
+                f, f"{test_set_name}-{key}", results, enable_log=True
+            )
+            test_set_wers[key] = wer
+
+        logging.info("Wrote detailed error stats to {}".format(errs_filename))
+
+    test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
+    errs_info = params.res_dir / f"wer-summary-{test_set_name}-{params.suffix}.txt"
+    with open(errs_info, "w") as f:
+        print("settings\tWER", file=f)
+        for key, val in test_set_wers:
+            print("{}\t{}".format(key, val), file=f)
+
+    s = "\nFor {}, WER of different settings are:\n".format(test_set_name)
+    note = "\tbest for {}".format(test_set_name)
+    for key, val in test_set_wers:
+        s += "{}\t{}{}\n".format(key, val, note)
+        note = ""
+    logging.info(s)
+
+
+@torch.no_grad()
+def main():
+    parser = get_parser()
+    AishellAsrDataModule.add_arguments(parser)
+    args = parser.parse_args()
+    args.exp_dir = Path(args.exp_dir)
+    args.lang_dir = Path(args.lang_dir)
+
+    params = get_params()
+    params.update(vars(args))
+
+    assert params.decoding_method in (
+        "greedy_search",
+        "beam_search",
+        "fast_beam_search",
+        "modified_beam_search",
+    )
+
+    if os.path.exists(params.context_file):
+        params.has_contexts = True
+    else:
+        params.has_contexts = False
+
+    params.res_dir = params.exp_dir / params.decoding_method
+
+    if params.iter > 0:
+        params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+    else:
+        params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+    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}"
+    elif "beam_search" in params.decoding_method:
+        params.suffix += f"-{params.decoding_method}-beam-size-{params.beam_size}"
+        if params.has_contexts:
+            params.suffix += f"-context-score-{params.context_score}"
+        else:
+            params.suffix += "-no-contexts-words"
+    else:
+        params.suffix += f"-context-{params.context_size}"
+        params.suffix += f"-max-sym-per-frame-{params.max_sym_per_frame}"
+    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}")
+
+    lexicon = Lexicon(params.lang_dir)
+    params.blank_id = 0
+    params.vocab_size = max(lexicon.tokens) + 1
+
+    graph_compiler = CharCtcTrainingGraphCompiler(
+        lexicon=lexicon,
+        device=device,
+    )
+
+    logging.info(params)
+
+    logging.info("About to create model")
+    model = get_transducer_model(params)
+    assert model.encoder.decode_chunk_size == params.decode_chunk_len // 2, (
+        model.encoder.decode_chunk_size,
+        params.decode_chunk_len,
+    )
+
+    if not params.use_averaged_model:
+        if 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()
+
+    if params.decoding_method == "fast_beam_search":
+        decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+    else:
+        decoding_graph = None
+
+    if params.decoding_method == "modified_beam_search":
+        if os.path.exists(params.context_file):
+            contexts_text = []
+            for line in open(params.context_file).readlines():
+                contexts_text.append(line.strip())
+            contexts = graph_compiler.texts_to_ids(contexts_text)
+            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}")
+
+    # we need cut ids to display recognition results.
+    args.return_cuts = True
+    aishell = AishellAsrDataModule(args)
+
+    test_cuts = aishell.test_cuts()
+    dev_cuts = aishell.valid_cuts()
+
+    test_dl = aishell.test_dataloaders(test_cuts)
+    dev_dl = aishell.test_dataloaders(dev_cuts)
+
+    test_sets = ["test", "dev"]
+    test_dls = [test_dl, dev_dl]
+
+    import time
+
+    for test_set, test_dl in zip(test_sets, test_dls):
+        start = time.time()
+        results_dict = decode_dataset(
+            dl=test_dl,
+            params=params,
+            model=model,
+            token_table=lexicon.token_table,
+            decoding_graph=decoding_graph,
+            context_graph=context_graph,
+        )
+        logging.info(f"Elasped time for {test_set}: {time.time() - start}")
+
+        save_results(
+            params=params,
+            test_set_name=test_set,
+            results_dict=results_dict,
+        )
+
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decode_stream.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decode_stream.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/decode_stream.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decoder.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/decoder.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/decoder.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/encoder_interface.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/encoder_interface.py
@ -0,0 +1 @@
+../pruned_transducer_stateless2/encoder_interface.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn-zh.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn-zh.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn-zh.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/export-for-ncnn.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-onnx-zh.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-onnx-zh.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx-zh.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/export-onnx.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/export.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/export.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_pretrained.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_pretrained.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/jit_pretrained.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_export.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/jit_trace_pretrained.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/joiner.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/joiner.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/joiner.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/model.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/model.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/model.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/ncnn_custom_layer.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/ncnn_custom_layer.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/ncnn_custom_layer.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_check.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_model_wrapper.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_model_wrapper.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_model_wrapper.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/onnx_pretrained.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/optim.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/optim.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/optim.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/pretrained.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/pretrained.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/pretrained.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/scaling.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/scaling.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/scaling_converter.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/scaling_converter.py
@ -0,0 +1 @@
+../pruned_transducer_stateless7/scaling_converter.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming-ncnn-decode.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming-ncnn-decode.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/streaming-ncnn-decode.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming_beam_search.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming_beam_search.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless2/streaming_beam_search.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming_decode.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/streaming_decode.py
@ -0,0 +1,627 @@
+#!/usr/bin/env python3
+# Copyright 2022 Xiaomi Corporation (Authors: Wei Kang, 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.
+
+"""
+Usage:
+./pruned_transducer_stateless7_streaming/streaming_decode.py \
+  --epoch 28 \
+  --avg 15 \
+  --decode-chunk-len 32 \
+  --exp-dir ./pruned_transducer_stateless7_streaming/exp \
+  --decoding-method greedy_search \
+  --num-decode-streams 2000
+"""
+
+import argparse
+import logging
+import math
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import k2
+import numpy as np
+import torch
+import torch.nn as nn
+from asr_datamodule import AishellAsrDataModule
+from decode_stream import DecodeStream
+from kaldifeat import Fbank, FbankOptions
+from lhotse import CutSet
+from streaming_beam_search import (
+    fast_beam_search_one_best,
+    greedy_search,
+    modified_beam_search,
+)
+from torch.nn.utils.rnn import pad_sequence
+from train import add_model_arguments, get_params, get_transducer_model
+from zipformer import stack_states, unstack_states
+
+from icefall import ContextGraph
+from icefall.char_graph_compiler import CharCtcTrainingGraphCompiler
+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 get_parser():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--epoch",
+        type=int,
+        default=28,
+        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(
+        "--exp-dir",
+        type=str,
+        default="pruned_transducer_stateless7_streaming/exp",
+        help="The experiment dir",
+    )
+
+    parser.add_argument(
+        "--decoding-method",
+        type=str,
+        default="greedy_search",
+        help="""Supported decoding methods are:
+        greedy_search
+        modified_beam_search
+        fast_beam_search
+        """,
+    )
+
+    parser.add_argument(
+        "--num_active_paths",
+        type=int,
+        default=4,
+        help="""An interger indicating how many candidates we will keep for each
+        frame. Used only when --decoding-method is modified_beam_search.""",
+    )
+
+    parser.add_argument(
+        "--beam",
+        type=float,
+        default=4,
+        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""",
+    )
+
+    parser.add_argument(
+        "--max-contexts",
+        type=int,
+        default=4,
+        help="""Used only when --decoding-method is
+        fast_beam_search""",
+    )
+
+    parser.add_argument(
+        "--max-states",
+        type=int,
+        default=32,
+        help="""Used only when --decoding-method is
+        fast_beam_search""",
+    )
+
+    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(
+        "--num-decode-streams",
+        type=int,
+        default=2000,
+        help="The number of streams that can be decoded parallel.",
+    )
+
+    add_model_arguments(parser)
+
+    return parser
+
+
+def decode_one_chunk(
+    params: AttributeDict,
+    model: nn.Module,
+    decode_streams: List[DecodeStream],
+) -> List[int]:
+    """Decode one chunk frames of features for each decode_streams and
+    return the indexes of finished streams in a List.
+
+    Args:
+      params:
+        It's the return value of :func:`get_params`.
+      model:
+        The neural model.
+      decode_streams:
+        A List of DecodeStream, each belonging to a utterance.
+    Returns:
+      Return a List containing which DecodeStreams are finished.
+    """
+    device = model.device
+
+    features = []
+    feature_lens = []
+    states = []
+    processed_lens = []
+
+    for stream in decode_streams:
+        feat, feat_len = stream.get_feature_frames(params.decode_chunk_len)
+        features.append(feat)
+        feature_lens.append(feat_len)
+        states.append(stream.states)
+        processed_lens.append(stream.done_frames)
+
+    feature_lens = torch.tensor(feature_lens, device=device)
+    features = pad_sequence(features, batch_first=True, padding_value=LOG_EPS)
+
+    # We subsample features with ((x_len - 7) // 2 + 1) // 2 and the max downsampling
+    # factor in encoders is 8.
+    # After feature embedding (x_len - 7) // 2, we have (23 - 7) // 2 = 8.
+    tail_length = 23
+    if features.size(1) < tail_length:
+        pad_length = tail_length - features.size(1)
+        feature_lens += pad_length
+        features = torch.nn.functional.pad(
+            features,
+            (0, 0, 0, pad_length),
+            mode="constant",
+            value=LOG_EPS,
+        )
+
+    states = stack_states(states)
+    processed_lens = torch.tensor(processed_lens, device=device)
+
+    encoder_out, encoder_out_lens, new_states = model.encoder.streaming_forward(
+        x=features,
+        x_lens=feature_lens,
+        states=states,
+    )
+
+    encoder_out = model.joiner.encoder_proj(encoder_out)
+
+    if params.decoding_method == "greedy_search":
+        greedy_search(model=model, encoder_out=encoder_out, streams=decode_streams)
+    elif params.decoding_method == "fast_beam_search":
+        processed_lens = processed_lens + encoder_out_lens
+        fast_beam_search_one_best(
+            model=model,
+            encoder_out=encoder_out,
+            processed_lens=processed_lens,
+            streams=decode_streams,
+            beam=params.beam,
+            max_states=params.max_states,
+            max_contexts=params.max_contexts,
+        )
+    elif params.decoding_method == "modified_beam_search":
+        modified_beam_search(
+            model=model,
+            streams=decode_streams,
+            encoder_out=encoder_out,
+            num_active_paths=params.num_active_paths,
+        )
+    else:
+        raise ValueError(f"Unsupported decoding method: {params.decoding_method}")
+
+    states = unstack_states(new_states)
+
+    finished_streams = []
+    for i in range(len(decode_streams)):
+        decode_streams[i].states = states[i]
+        decode_streams[i].done_frames += encoder_out_lens[i]
+        if decode_streams[i].done:
+            finished_streams.append(i)
+
+    return finished_streams
+
+
+def decode_dataset(
+    cuts: CutSet,
+    params: AttributeDict,
+    model: nn.Module,
+    token_table: Optional[k2.SymbolTable] = None,
+    decoding_graph: Optional[k2.Fsa] = None,
+    context_graph: Optional[ContextGraph] = None,
+) -> Dict[str, List[Tuple[List[str], List[str]]]]:
+    """Decode dataset.
+
+    Args:
+      cuts:
+        Lhotse Cutset containing the dataset to decode.
+      params:
+        It is returned by :func:`get_params`.
+      model:
+        The neural model.
+      sp:
+        The BPE model.
+      decoding_graph:
+        The decoding graph. Can be either a `k2.trivial_graph` or HLG, Used
+        only when --decoding_method is fast_beam_search.
+    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.
+    """
+    device = model.device
+
+    opts = FbankOptions()
+    opts.device = device
+    opts.frame_opts.dither = 0
+    opts.frame_opts.snip_edges = False
+    opts.frame_opts.samp_freq = 16000
+    opts.mel_opts.num_bins = 80
+
+    log_interval = 50
+
+    decode_results = []
+    # Contain decode streams currently running.
+    decode_streams = []
+    for num, cut in enumerate(cuts):
+        # each utterance has a DecodeStream.
+        initial_states = model.encoder.get_init_state(device=device)
+        decode_stream = DecodeStream(
+            params=params,
+            cut_id=cut.id,
+            initial_states=initial_states,
+            decoding_graph=decoding_graph,
+            device=device,
+        )
+
+        audio: np.ndarray = cut.load_audio()
+        # audio.shape: (1, num_samples)
+        assert len(audio.shape) == 2
+        assert audio.shape[0] == 1, "Should be single channel"
+        assert audio.dtype == np.float32, audio.dtype
+
+        # The trained model is using normalized samples
+        assert audio.max() <= 1, "Should be normalized to [-1, 1])"
+
+        samples = torch.from_numpy(audio).squeeze(0)
+
+        fbank = Fbank(opts)
+        feature = fbank(samples.to(device))
+        decode_stream.set_features(feature, tail_pad_len=params.decode_chunk_len)
+        decode_stream.ground_truth = cut.supervisions[0].text
+
+        decode_streams.append(decode_stream)
+
+        while len(decode_streams) >= params.num_decode_streams:
+            finished_streams = decode_one_chunk(
+                params=params, model=model, decode_streams=decode_streams
+            )
+            for i in sorted(finished_streams, reverse=True):
+                decode_results.append(
+                    (
+                        decode_streams[i].id,
+                        decode_streams[i].ground_truth.split(),
+                        [
+                            token_table[result]
+                            for result in decode_streams[i].decoding_result()
+                        ],
+                    )
+                )
+                del decode_streams[i]
+
+        if num % log_interval == 0:
+            logging.info(f"Cuts processed until now is {num}.")
+
+    # decode final chunks of last sequences
+    while len(decode_streams):
+        finished_streams = decode_one_chunk(
+            params=params, model=model, decode_streams=decode_streams
+        )
+        for i in sorted(finished_streams, reverse=True):
+            decode_results.append(
+                (
+                    decode_streams[i].id,
+                    decode_streams[i].ground_truth.split(),
+                    [
+                        token_table[result]
+                        for result in decode_streams[i].decoding_result()
+                    ],
+                )
+            )
+            del decode_streams[i]
+
+    if params.decoding_method == "greedy_search":
+        key = "greedy_search"
+    elif params.decoding_method == "fast_beam_search":
+        key = (
+            f"beam_{params.beam}_"
+            f"max_contexts_{params.max_contexts}_"
+            f"max_states_{params.max_states}"
+        )
+    elif params.decoding_method == "modified_beam_search":
+        key = f"num_active_paths_{params.num_active_paths}"
+    else:
+        raise ValueError(f"Unsupported decoding method: {params.decoding_method}")
+    return {key: decode_results}
+
+
+def save_results(
+    params: AttributeDict,
+    test_set_name: str,
+    results_dict: Dict[str, List[Tuple[List[str], List[str]]]],
+):
+    test_set_wers = dict()
+    for key, results in results_dict.items():
+        recog_path = params.res_dir / f"recogs-{test_set_name}-{params.suffix}.txt"
+        results = sorted(results)
+        store_transcripts(filename=recog_path, texts=results)
+        logging.info(f"The transcripts are stored in {recog_path}")
+
+        # The following prints out WERs, per-word error statistics and aligned
+        # ref/hyp pairs.
+        errs_filename = params.res_dir / f"errs-{test_set_name}-{params.suffix}.txt"
+        with open(errs_filename, "w") as f:
+            wer = write_error_stats(
+                f, f"{test_set_name}-{key}", results, enable_log=True
+            )
+            test_set_wers[key] = wer
+
+        logging.info("Wrote detailed error stats to {}".format(errs_filename))
+
+    test_set_wers = sorted(test_set_wers.items(), key=lambda x: x[1])
+    errs_info = params.res_dir / f"wer-summary-{test_set_name}-{params.suffix}.txt"
+    with open(errs_info, "w") as f:
+        print("settings\tWER", file=f)
+        for key, val in test_set_wers:
+            print("{}\t{}".format(key, val), file=f)
+
+    s = "\nFor {}, WER of different settings are:\n".format(test_set_name)
+    note = "\tbest for {}".format(test_set_name)
+    for key, val in test_set_wers:
+        s += "{}\t{}{}\n".format(key, val, note)
+        note = ""
+    logging.info(s)
+
+
+@torch.no_grad()
+def main():
+    parser = get_parser()
+    AishellAsrDataModule.add_arguments(parser)
+    args = parser.parse_args()
+    args.exp_dir = Path(args.exp_dir)
+
+    params = get_params()
+    params.update(vars(args))
+
+    params.res_dir = params.exp_dir / "streaming" / params.decoding_method
+
+    if params.iter > 0:
+        params.suffix = f"iter-{params.iter}-avg-{params.avg}"
+    else:
+        params.suffix = f"epoch-{params.epoch}-avg-{params.avg}"
+
+    # for streaming
+    params.suffix += f"-streaming-chunk-size-{params.decode_chunk_len}"
+
+    # for fast_beam_search
+    if params.decoding_method == "fast_beam_search":
+        params.suffix += f"-beam-{params.beam}"
+        params.suffix += f"-max-contexts-{params.max_contexts}"
+        params.suffix += f"-max-states-{params.max_states}"
+
+    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}")
+
+    lexicon = Lexicon(params.lang_dir)
+    params.blank_id = 0
+    params.vocab_size = max(lexicon.tokens) + 1
+
+    graph_compiler = CharCtcTrainingGraphCompiler(
+        lexicon=lexicon,
+        device=device,
+    )
+
+    logging.info(params)
+
+    logging.info("About to create model")
+    model = get_transducer_model(params)
+
+    if not params.use_averaged_model:
+        if 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 start >= 0:
+                    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()
+    model.device = device
+
+    decoding_graph = None
+    if params.decoding_method == "fast_beam_search":
+        decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
+
+    if params.decoding_method == "modified_beam_search":
+        if os.path.exists(params.context_file):
+            contexts_text = []
+            for line in open(params.context_file).readlines():
+                contexts_text.append(line.strip())
+            contexts = graph_compiler.texts_to_ids(contexts_text)
+            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}")
+
+    aishell = AishellAsrDataModule(args)
+
+    test_cuts = aishell.test_cuts()
+    valid_cuts = aishell.valid_cuts()
+
+    test_sets = ["test", "valid"]
+    cuts = [test_cuts, valid_cuts]
+
+    for test_set, test_cut in zip(test_sets, cuts):
+        results_dict = decode_dataset(
+            cuts=test_cut,
+            params=params,
+            model=model,
+            token_table=lexicon.token_table,
+            decoding_graph=decoding_graph,
+            context_graph=context_graph,
+        )
+
+        save_results(
+            params=params,
+            test_set_name=test_set,
+            results_dict=results_dict,
+        )
+
+    logging.info("Done!")
+
+
+if __name__ == "__main__":
+    main()
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/test_model.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/test_model.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/test_model.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/train.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/train.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/train2.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/train2.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/zipformer.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/zipformer.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer.py
--- a/egs/aishell/ASR/pruned_transducer_stateless7_streaming/zipformer2.py
+++ b/egs/aishell/ASR/pruned_transducer_stateless7_streaming/zipformer2.py
@ -0,0 +1 @@
+../../../librispeech/ASR/pruned_transducer_stateless7_streaming/zipformer2.py
--- a/egs/aishell/ASR/tdnn_lstm_ctc/asr_datamodule.py
+++ b/egs/aishell/ASR/tdnn_lstm_ctc/asr_datamodule.py
@ -198,7 +198,7 @@ class AishellAsrDataModule:
        if self.args.enable_musan:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/aishell/ASR/transducer_stateless_modified-2/asr_datamodule.py
+++ b/egs/aishell/ASR/transducer_stateless_modified-2/asr_datamodule.py
@ -160,7 +160,7 @@ class AsrDataModule:
        if cuts_musan is not None:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/aishell2/ASR/pruned_transducer_stateless5/asr_datamodule.py
+++ b/egs/aishell2/ASR/pruned_transducer_stateless5/asr_datamodule.py
@ -218,7 +218,7 @@ class AiShell2AsrDataModule:
            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, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/aishell2/ASR/pruned_transducer_stateless5/train.py
+++ b/egs/aishell2/ASR/pruned_transducer_stateless5/train.py
@ -730,7 +730,6 @@ def train_one_epoch(
    tot_loss = MetricsTracker()

    for batch_idx, batch in enumerate(train_dl):
-
        params.batch_idx_train += 1
        batch_size = len(batch["supervisions"]["text"])

@ -919,7 +918,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/aishell4/ASR/pruned_transducer_stateless5/asr_datamodule.py
+++ b/egs/aishell4/ASR/pruned_transducer_stateless5/asr_datamodule.py
@ -228,7 +228,7 @@ class Aishell4AsrDataModule:
        if self.args.enable_musan:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/aishell4/ASR/pruned_transducer_stateless5/train.py
+++ b/egs/aishell4/ASR/pruned_transducer_stateless5/train.py
@ -908,7 +908,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/alimeeting/ASR/pruned_transducer_stateless2/asr_datamodule.py
+++ b/egs/alimeeting/ASR/pruned_transducer_stateless2/asr_datamodule.py
@ -211,7 +211,7 @@ class AlimeetingAsrDataModule:
        if self.args.enable_musan:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/alimeeting/ASR/pruned_transducer_stateless2/train.py
+++ b/egs/alimeeting/ASR/pruned_transducer_stateless2/train.py
@ -635,7 +635,6 @@ def train_one_epoch(
    tot_loss = MetricsTracker()

    for batch_idx, batch in enumerate(train_dl):
-
        params.batch_idx_train += 1
        batch_size = len(batch["supervisions"]["text"])

@ -800,7 +799,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/asr_datamodule.py
@ -208,7 +208,7 @@ class AlimeetingAsrDataModule:
            logging.info("Enable MUSAN")
            cuts_musan = load_manifest(self.args.manifest_dir / "musan_cuts.jsonl.gz")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
@ -288,7 +288,6 @@ class AlimeetingAsrDataModule:
        return train_dl

    def valid_dataloaders(self, cuts_valid: CutSet) -> DataLoader:
-
        transforms = []
        if self.args.concatenate_cuts:
            transforms = [
--- a/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/train.py
+++ b/egs/alimeeting/ASR_v2/pruned_transducer_stateless7/train.py
@ -999,7 +999,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/ami/ASR/pruned_transducer_stateless7/asr_datamodule.py
+++ b/egs/ami/ASR/pruned_transducer_stateless7/asr_datamodule.py
@ -214,7 +214,7 @@ class AmiAsrDataModule:
            logging.info("Enable MUSAN")
            cuts_musan = load_manifest(self.args.manifest_dir / "musan_cuts.jsonl.gz")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/ami/ASR/pruned_transducer_stateless7/train.py
+++ b/egs/ami/ASR/pruned_transducer_stateless7/train.py
@ -988,7 +988,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/ami/SURT/dprnn_zipformer/asr_datamodule.py
+++ b/egs/ami/SURT/dprnn_zipformer/asr_datamodule.py
@ -202,7 +202,7 @@ class AmiAsrDataModule:
            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, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/commonvoice/ASR/pruned_transducer_stateless7/asr_datamodule.py
+++ b/egs/commonvoice/ASR/pruned_transducer_stateless7/asr_datamodule.py
@ -230,7 +230,7 @@ class CommonVoiceAsrDataModule:
            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, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/commonvoice/ASR/pruned_transducer_stateless7/train.py
+++ b/egs/commonvoice/ASR/pruned_transducer_stateless7/train.py
@ -1019,7 +1019,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/csj/ASR/local/utils/asr_datamodule.py
+++ b/egs/csj/ASR/local/utils/asr_datamodule.py
@ -256,7 +256,7 @@ class CSJAsrDataModule:
            logging.info("About to get Musan cuts")
            cuts_musan = load_manifest(self.args.musan_dir / "musan_cuts.jsonl.gz")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/csj/ASR/pruned_transducer_stateless7_streaming/train.py
+++ b/egs/csj/ASR/pruned_transducer_stateless7_streaming/train.py
@ -1074,7 +1074,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/csj/ASR/pruned_transducer_stateless7_streaming/train2.py
+++ b/egs/csj/ASR/pruned_transducer_stateless7_streaming/train2.py
@ -1075,7 +1075,7 @@ def run(rank, world_size, args):

    if params.print_diagnostics:
        opts = diagnostics.TensorDiagnosticOptions(
-            2**22
+            512
        )  # allow 4 megabytes per sub-module
        diagnostic = diagnostics.attach_diagnostics(model, opts)

--- a/egs/gigaspeech/ASR/README.md
+++ b/egs/gigaspeech/ASR/README.md
@ -15,6 +15,7 @@ ln -sfv /path/to/GigaSpeech download/GigaSpeech
 ## Performance Record
 |                                |  Dev  | Test  |
 |--------------------------------|-------|-------|
+|           `zipformer`          | 10.25 | 10.38 |
 |         `conformer_ctc`        | 10.47 | 10.58 |
 | `pruned_transducer_stateless2` | 10.40 | 10.51 |

--- a/egs/gigaspeech/ASR/RESULTS.md
+++ b/egs/gigaspeech/ASR/RESULTS.md
@ -1,4 +1,78 @@
 ## Results
+### zipformer (zipformer + pruned stateless transducer)
+
+See <https://github.com/k2-fsa/icefall/pull/1254> for more details.
+
+[zipformer](./zipformer)
+
+- Non-streaming
+- normal-scaled model, number of model parameters: 65549011, i.e., 65.55 M
+
+You can find a pretrained model, training logs, decoding logs, and decoding results at:
+<https://huggingface.co/yfyeung/icefall-asr-gigaspeech-zipformer-2023-10-17>
+
+The tensorboard log for training is available at
+<https://wandb.ai/yifanyeung/icefall-asr-gigaspeech-zipformer-2023-10-20>
+
+You can use <https://github.com/k2-fsa/sherpa> to deploy it.
+
+| decoding method      | test-clean | test-other | comment            |
+|----------------------|------------|------------|--------------------|
+| greedy_search        | 10.31      | 10.50      | --epoch 30 --avg 9 |
+| modified_beam_search | 10.25      | 10.38      | --epoch 30 --avg 9 |
+| fast_beam_search     | 10.26      | 10.48      | --epoch 30 --avg 9 |
+
+The training command is:
+```bash
+export CUDA_VISIBLE_DEVICES="0,1,2,3"
+./zipformer/train.py \
+  --world-size 4 \
+  --num-epochs 30 \
+  --start-epoch 1 \
+  --use-fp16 1 \
+  --exp-dir zipformer/exp \
+  --causal 0 \
+  --subset XL \
+  --max-duration 700 \
+  --use-transducer 1 \
+  --use-ctc 0 \
+  --lr-epochs 1 \
+  --master-port 12345
+```
+
+The decoding command is:
+```bash
+export CUDA_VISIBLE_DEVICES=0
+
+# greedy search
+./zipformer/decode.py \
+  --epoch 30 \
+  --avg 9 \
+  --exp-dir ./zipformer/exp \
+  --max-duration 1000 \
+  --decoding-method greedy_search
+
+# modified beam search
+./zipformer/decode.py \
+  --epoch 30 \
+  --avg 9 \
+  --exp-dir ./zipformer/exp \
+  --max-duration 1000 \
+  --decoding-method modified_beam_search \
+  --beam-size 4
+
+# fast beam search (one best)
+./zipformer/decode.py \
+  --epoch 30 \
+  --avg 9 \
+  --exp-dir ./zipformer/exp \
+  --max-duration 1000 \
+  --decoding-method fast_beam_search \
+  --beam 20.0 \
+  --max-contexts 8 \
+  --max-states 64
+```
+
 ### GigaSpeech BPE training results (Pruned Transducer 2)

 #### 2022-05-12
--- a/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py
+++ b/egs/gigaspeech/ASR/conformer_ctc/asr_datamodule.py
@ -194,7 +194,7 @@ class GigaSpeechAsrDataModule:
        if self.args.enable_musan:
            logging.info("Enable MUSAN")
            transforms.append(
-                CutMix(cuts=cuts_musan, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py
+++ b/egs/gigaspeech/ASR/pruned_transducer_stateless2/asr_datamodule.py
@ -217,7 +217,7 @@ class GigaSpeechAsrDataModule:
            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, prob=0.5, snr=(10, 20), preserve_id=True)
+                CutMix(cuts=cuts_musan, p=0.5, snr=(10, 20), preserve_id=True)
            )
        else:
            logging.info("Disable MUSAN")
--- a/Show More
+++ b/Show More
				`@ -0,0 +1 @@`
				`../../../librispeech/ASR/conformer_ctc/jit_pretrained_decode_with_H.py`
				`@ -0,0 +1 @@`
				`../../../librispeech/ASR/local/prepare_lang_fst.py`
				`@ -0,0 +1 @@`
				`../../../librispeech/ASR/pruned_transducer_stateless7_streaming/README.md`
				`@ -0,0 +1 @@`
				`../pruned_transducer_stateless2/asr_datamodule.py`
				`@ -0,0 +1 @@`
				`../pruned_transducer_stateless2/beam_search.py`
				`@ -0,0 +1 @@`
				`../pruned_transducer_stateless2/encoder_interface.py`
				`@ -0,0 +1 @@`
				`../pruned_transducer_stateless7/scaling_converter.py`