update export models and scripts (#6)

- Use the latest pruned_transducer_stateless7_streaming script in icefall to export the model and update the export scripts in the repo (427c4b37e1885d621595389453cad0717fbad7bd)

exp/cpu_jit.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5fe5d637dc8f8a8717c4cad0200909271ef4b30248542896f0257a726c289631
-size 379190270
+oid sha256:3b41fa49583b69438105016d68672604e0359498925dd0c7b5965184a445cc8c
+size 379196926

exp/decoder_jit_trace.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f6298c4c09b88b0020b4fab955ff9dea41a00bbf931db2007f6a70b412346d33
-size 12831269
+oid sha256:83cc6f5cbf4e3e7a518546c2ee4e8d9c17d479ceebcac3648a031985d44ec89d
+size 12831333

exp/encoder_jit_trace.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d748e243b4dd566196c9488346ff3863cad04d6ade6f3336c81491a5b5f0f9f0
-size 330595617
+oid sha256:83e2ec17083607da4085d9653c7d15c908e28fd50488e62c6f24e11e795a90a4
+size 330607841

exp/export-stateless7-streaming-zh.sh

@@ -2,8 +2,8 @@
 
 . path.sh
 
-./pruned_transducer_stateless7_streaming/export-zh.py \
-  --lang-dir ./k2fsa-zipformer-chinese-english-mixed/data/lang_char_bpe \
+./pruned_transducer_stateless7_streaming/export.py \
+  --tokens ./k2fsa-zipformer-chinese-english-mixed/data/lang_char_bpe/tokens.txt \
   --use-averaged-model 0 \
   --epoch 99 \
   --avg 1 \

exp/jit_trace_export-zh.py

@@ -1,323 +0,0 @@
-#!/usr/bin/env python3
-
-"""
-Usage:
-./pruned_transducer_stateless7_streaming/jit_trace_export-zh.py \
-  --exp-dir $dir/exp \
-  --exp-dir ./pruned_transducer_stateless7_streaming/exp \
-  --lang-dir ./data/lang_char_bpe \
-  --epoch 99 \
-  --avg 1 \
-  --use-averaged-model 0 \
-  \
-  --decode-chunk-len 32 \
-  --num-encoder-layers "2,4,3,2,4" \
-  --feedforward-dims "1024,1024,1536,1536,1024" \
-  --nhead "8,8,8,8,8" \
-  --encoder-dims "384,384,384,384,384" \
-  --attention-dims "192,192,192,192,192" \
-  --encoder-unmasked-dims "256,256,256,256,256" \
-  --zipformer-downsampling-factors "1,2,4,8,2" \
-  --cnn-module-kernels "31,31,31,31,31" \
-  --decoder-dim 512 \
-  --joiner-dim 512
-"""
-
-import argparse
-import logging
-from pathlib import Path
-
-import sentencepiece as spm
-import torch
-from scaling_converter import convert_scaled_to_non_scaled
-from train import add_model_arguments, get_params, get_transducer_model
-from icefall.lexicon import Lexicon
-
-from icefall.checkpoint import (
-    average_checkpoints,
-    average_checkpoints_with_averaged_model,
-    find_checkpoints,
-    load_checkpoint,
-)
-from icefall.utils import AttributeDict, str2bool
-
-
-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 averaging.
-        Note: Epoch counts from 0.
-        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(
-        "--exp-dir",
-        type=str,
-        default="pruned_transducer_stateless2/exp",
-        help="""It specifies the directory where all training related
-        files, e.g., checkpoints, log, etc, are saved
-        """,
-    )
-
-    parser.add_argument(
-        "--lang-dir",
-        type=str,
-        default="data/lang_char",
-        help="The lang dir",
-    )
-
-    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(
-        "--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. ",
-    )
-
-    add_model_arguments(parser)
-
-    return parser
-
-
-def export_encoder_model_jit_trace(
-    encoder_model: torch.nn.Module,
-    encoder_filename: str,
-    params: AttributeDict,
-) -> None:
-    """Export the given encoder model with torch.jit.trace()
-
-    Note: The warmup argument is fixed to 1.
-
-    Args:
-      encoder_model:
-        The input encoder model
-      encoder_filename:
-        The filename to save the exported model.
-    """
-    decode_chunk_len = params.decode_chunk_len  # before subsampling
-    pad_length = 7
-    s = f"decode_chunk_len: {decode_chunk_len}"
-    logging.info(s)
-    assert encoder_model.decode_chunk_size == decode_chunk_len // 2, (
-        encoder_model.decode_chunk_size,
-        decode_chunk_len,
-    )
-
-    T = decode_chunk_len + pad_length
-
-    x = torch.zeros(1, T, 80, dtype=torch.float32)
-    x_lens = torch.full((1,), T, dtype=torch.int32)
-    states = encoder_model.get_init_state(device=x.device)
-
-    encoder_model.__class__.forward = encoder_model.__class__.streaming_forward
-    traced_model = torch.jit.trace(encoder_model, (x, x_lens, states))
-    traced_model.save(encoder_filename)
-    logging.info(f"Saved to {encoder_filename}")
-
-
-def export_decoder_model_jit_trace(
-    decoder_model: torch.nn.Module,
-    decoder_filename: str,
-) -> None:
-    """Export the given decoder model with torch.jit.trace()
-
-    Note: The argument need_pad is fixed to False.
-
-    Args:
-      decoder_model:
-        The input decoder model
-      decoder_filename:
-        The filename to save the exported model.
-    """
-    y = torch.zeros(10, decoder_model.context_size, dtype=torch.int64)
-    need_pad = torch.tensor([False])
-
-    traced_model = torch.jit.trace(decoder_model, (y, need_pad))
-    traced_model.save(decoder_filename)
-    logging.info(f"Saved to {decoder_filename}")
-
-
-def export_joiner_model_jit_trace(
-    joiner_model: torch.nn.Module,
-    joiner_filename: str,
-) -> None:
-    """Export the given joiner model with torch.jit.trace()
-
-    Note: The argument project_input is fixed to True. A user should not
-    project the encoder_out/decoder_out by himself/herself. The exported joiner
-    will do that for the user.
-
-    Args:
-      joiner_model:
-        The input joiner model
-      joiner_filename:
-        The filename to save the exported model.
-
-    """
-    encoder_out_dim = joiner_model.encoder_proj.weight.shape[1]
-    decoder_out_dim = joiner_model.decoder_proj.weight.shape[1]
-    encoder_out = torch.rand(1, encoder_out_dim, dtype=torch.float32)
-    decoder_out = torch.rand(1, decoder_out_dim, dtype=torch.float32)
-
-    traced_model = torch.jit.trace(joiner_model, (encoder_out, decoder_out))
-    traced_model.save(joiner_filename)
-    logging.info(f"Saved to {joiner_filename}")
-
-
-@torch.no_grad()
-def main():
-    args = get_parser().parse_args()
-    args.exp_dir = Path(args.exp_dir)
-
-    params = get_params()
-    params.update(vars(args))
-
-    device = torch.device("cpu")
-
-    logging.info(f"device: {device}")
-
-    lexicon = Lexicon(params.lang_dir)
-    params.blank_id = 0
-    params.vocab_size = max(lexicon.tokens) + 1
-
-    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 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("cpu")
-    model.eval()
-
-    convert_scaled_to_non_scaled(model, inplace=True)
-    logging.info("Using torch.jit.trace()")
-
-    logging.info("Exporting encoder")
-    encoder_filename = params.exp_dir / "encoder_jit_trace.pt"
-    export_encoder_model_jit_trace(model.encoder, encoder_filename, params)
-
-    logging.info("Exporting decoder")
-    decoder_filename = params.exp_dir / "decoder_jit_trace.pt"
-    export_decoder_model_jit_trace(model.decoder, decoder_filename)
-
-    logging.info("Exporting joiner")
-    joiner_filename = params.exp_dir / "joiner_jit_trace.pt"
-    export_joiner_model_jit_trace(model.joiner, joiner_filename)
-
-
-if __name__ == "__main__":
-    formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
-
-    logging.basicConfig(format=formatter, level=logging.INFO)
-    main()

exp/jit_trace_export-zh.sh

@@ -15,9 +15,9 @@ if [ ! -f $dir/exp/epoch-99.pt ]; then
   popd
 fi
 
-./pruned_transducer_stateless7_streaming/jit_trace_export-zh.py \
+./pruned_transducer_stateless7_streaming/jit_trace_export.py \
   --exp-dir $dir/exp \
-  --lang-dir $dir/data/lang_char_bpe \
+  --bpe-model $dir/data/lang_char_bpe/bpe.model \
   --epoch 99 \
   --avg 1 \
   --use-averaged-model 0 \

exp/joiner_jit_trace.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1765e66b163ba43b0d1825e8ecb3f2edc23170ba5104b7e3827bf0bad318b12e
-size 14680987
+oid sha256:bdd58624fd2df70b5583c684c6705cd28b27d73bc15ef95484356315d1579043
+size 14681115