From 5f6b8056c9c13239c652b2acb7cd5087d7c40646 Mon Sep 17 00:00:00 2001
From: Yifan Yang <yifanyeung@qq.com>
Date: Fri, 2 Jun 2023 11:45:41 +0800
Subject: [PATCH] add export.py

---
 icefall/rnn_lm/export.py | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/icefall/rnn_lm/export.py b/icefall/rnn_lm/export.py
index 78a914cf4..8657ea9f2 100644
--- a/icefall/rnn_lm/export.py
+++ b/icefall/rnn_lm/export.py
@@ -19,6 +19,39 @@
 
 # This script converts several saved checkpoints
 # to a single one using model averaging.
+"""
+
+Usage:
+
+(1) Export to torchscript model using torch.jit.script()
+
+./rnn_lm/export.py \
+  --exp-dir ./rnn_lm/exp \
+  --epoch 20 \
+  --avg 10 \
+  --jit 1
+
+It will generate a file `cpu_jit.pt` in the given `exp_dir`. You can later
+load it by `torch.jit.load("cpu_jit.pt")`.
+
+Note `cpu` in the name `cpu_jit.pt` means the parameters when loaded into Python
+are on CPU. You can use `to("cuda")` to move them to a CUDA device.
+
+Check
+https://github.com/k2-fsa/sherpa
+for how to use the exported models outside of icefall.
+
+(2) Export `model.state_dict()`
+
+./rnn_lm/export.py \
+  --exp-dir ./rnn_lm/exp \
+  --bpe-model data/lang_bpe_500/bpe.model \
+  --epoch 20 \
+  --avg 10
+
+It will generate a file `pretrained.pt` in the given `exp_dir`. You can later
+load it by `icefall.checkpoint.load_checkpoint()`.
+"""
 
 import argparse
 import logging
@@ -83,7 +116,7 @@ def get_parser():
     parser.add_argument(
         "--exp-dir",
         type=str,
-        default="rnnlm/exp",
+        default="rnn_lm/exp",
         help="""It specifies the directory where all training related
         files, e.g., checkpoints, log, etc, are saved
         """,