#!/usr/bin/env bash

# Please go to
# https://huggingface.co/csukuangfj/icefall-asr-librispeech-lstm-transducer-stateless2-2022-09-03
# to download the pre-trained models

if [ ! -d lstm2 ]; then
  mkdir lstm2
  pushd lstm2
  ln -s ~/open-source/icefall-models/icefall-asr-librispeech-lstm-transducer-stateless2-2022-09-03/exp/pretrained-iter-468000-avg-16.pt epoch-99.pt
  ln -s ~/open-source/icefall-models/icefall-asr-librispeech-lstm-transducer-stateless2-2022-09-03/data .
  popd
fi

export CUDA_VISIBLE_DEVICES=""
set -ex

./lstm_transducer_stateless2/export-for-ncnn.py \
  --exp-dir ./lstm2 \
  --bpe-model ./lstm2/data/lang_bpe_500/bpe.model \
  --epoch 99 \
  --avg 1 \
  --use-averaged-model 0 \
  --num-encoder-layers 12 \
  --encoder-dim 512 \
  --rnn-hidden-size 1024

cd ./lstm2

pnnx encoder_jit_trace-pnnx.pt
pnnx decoder_jit_trace-pnnx.pt
pnnx joiner_jit_trace-pnnx.pt

# Now we need to modified encoder_jit_trace-pnnx.ncnn.param for sherpa-ncnn
# The following is the diff
#
# diff -uN ./encoder_jit_trace-pnnx.ncnn.param-before ./encoder_jit_trace-pnnx.ncnn.param
# --- ./encoder_jit_trace-pnnx.ncnn.param-before  2023-02-15 11:27:14.000000000 +0800
# +++ ./encoder_jit_trace-pnnx.ncnn.param 2023-02-15 11:35:41.000000000 +0800
# @@ -1,5 +1,6 @@
#  7767517
# -267 379
# +268 379
# +SherpaMetaData           sherpa_meta_data1        0 0 0=3 1=12 2=512 3=1024
#  Input                    in0                      0 1 in0
#  Input                    in1                      0 1 in1
#  Input                    in2                      0 1 in2

#
#
# Explanation:
# 1. A new layer is added, so we change the layer number 267 to 268
# 2. SherpaMetaData, it is the layer type. MUST be SherpaMetaData
# 3. sherpa_meta_data1, it is the layer name. MUST be sherpa_meta_data1
# 4. 0=3, attribute 0, MUST be 3
# 5. 1=12, attribute 1, 12 is the value of --num-encoder-layers
# 6. 2=512, attribute 2, 512 is the value of --encoder-dim
# 7. 3=1024, attribute 3, 1024 is the value of --rnn-hidden-size