text: #!/usr/bin/env python3
  #
  # Copyright 2021 Xiaomi Corporation (Author: 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.
  
  # This script converts several saved checkpoints
  # to a single one using model averaging.
  """
  Usage:
  ./transducer_stateless2/export.py \
    --exp-dir ./transducer_stateless2/exp \
    --bpe-model data/lang_bpe_500/bpe.model \
    --epoch 20 \
    --avg 10
  
  It will generate a file exp_dir/pretrained.pt
  
  To use the generated file with `transducer_stateless2/decode.py`, you can do:
  
      cd /path/to/exp_dir
      ln -s pretrained.pt epoch-9999.pt
  
      cd /path/to/egs/librispeech/ASR
      ./transducer_stateless2/decode.py \
          --exp-dir ./transducer_stateless2/exp \
          --epoch 9999 \
          --avg 1 \
          --max-duration 1 \
          --bpe-model data/lang_bpe_500/bpe.model
  """
  
  import argparse
  import logging
  from pathlib import Path
  
  import sentencepiece as spm
  import torch
  from train import get_params, get_transducer_model
  
  from icefall.checkpoint import average_checkpoints, load_checkpoint
  from icefall.utils import str2bool
  
  
  def get_parser():
      parser = argparse.ArgumentParser(
          formatter_class=argparse.ArgumentDefaultsHelpFormatter
      )
  
      parser.add_argument(
          "--epoch",
          type=int,
          default=20,
          help="It specifies the checkpoint to use for decoding."
          "Note: Epoch counts from 0.",
      )
  
      parser.add_argument(
          "--avg",
          type=int,
          default=10,
          help="Number of checkpoints to average. Automatically select "
          "consecutive checkpoints before the checkpoint specified by "
          "'--epoch'. ",
      )
  
      parser.add_argument(
          "--exp-dir",
          type=str,
          default="transducer_stateless2/exp",
          help="""It specifies the directory where all training related
          files, e.g., checkpoints, log, etc, are saved
          """,
      )
  
      parser.add_argument(
          "--bpe-model",
          type=str,
          default="data/lang_bpe_500/bpe.model",
          help="Path to the BPE model",
      )
  
      parser.add_argument(
          "--jit",
          type=str2bool,
          default=False,
          help="""True to save a model after applying torch.jit.script.
          """,
      )
  
      parser.add_argument(
          "--context-size",
          type=int,
          default=2,
          help="The context size in the decoder. 1 means bigram; "
          "2 means tri-gram",
      )
  
      return parser
  
  
  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")
      if torch.cuda.is_available():
          device = torch.device("cuda", 0)
  
      logging.info(f"device: {device}")
  
      sp = spm.SentencePieceProcessor()
      sp.load(params.bpe_model)
  
      # <blk> is defined in local/train_bpe_model.py
      params.blank_id = sp.piece_to_id("<blk>")
      params.vocab_size = sp.get_piece_size()
  
      logging.info(params)
  
      logging.info("About to create model")
      model = get_transducer_model(params)
  
      model.to(device)
  
      if 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))
  
      model.eval()
  
      model.to("cpu")
      model.eval()
  
      if params.jit:
          # We won't use the forward() method of the model in C++, so just ignore
          # it here.
          # Otherwise, one of its arguments is a ragged tensor and is not
          # torch scriptabe.
          model.__class__.forward = torch.jit.ignore(model.__class__.forward)
          logging.info("Using torch.jit.script")
          model = torch.jit.script(model)
          filename = params.exp_dir / "cpu_jit.pt"
          model.save(str(filename))
          logging.info(f"Saved to {filename}")
      else:
          logging.info("Not using torch.jit.script")
          # Save it using a format so that it can be loaded
          # by :func:`load_checkpoint`
          filename = params.exp_dir / "pretrained.pt"
          torch.save({"model": model.state_dict()}, str(filename))
          logging.info(f"Saved to {filename}")
  
  
  if __name__ == "__main__":
      formatter = (
          "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s"
      )
  
      logging.basicConfig(format=formatter, level=logging.INFO)
      main()
  