From 4d4a2c08f75242ecaea5d0bd3a6c60dadc336e30 Mon Sep 17 00:00:00 2001 From: marcoyang Date: Mon, 13 Feb 2023 12:04:29 +0800 Subject: [PATCH] fix export for ncnn --- .../lstm_transducer_stateless3/export-for-ncnn.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/egs/tal_csasr/ASR/lstm_transducer_stateless3/export-for-ncnn.py b/egs/tal_csasr/ASR/lstm_transducer_stateless3/export-for-ncnn.py index 7982ace68..d7bd3d9d9 100755 --- a/egs/tal_csasr/ASR/lstm_transducer_stateless3/export-for-ncnn.py +++ b/egs/tal_csasr/ASR/lstm_transducer_stateless3/export-for-ncnn.py @@ -60,6 +60,7 @@ from icefall.checkpoint import ( find_checkpoints, load_checkpoint, ) +from icefall.lexicon import Lexicon from icefall.utils import setup_logger, str2bool @@ -106,10 +107,10 @@ def get_parser(): ) parser.add_argument( - "--bpe-model", + "--lang-dir", type=str, - default="data/lang_bpe_500/bpe.model", - help="Path to the BPE model", + default="data/lang_char", + help="Path to the lang", ) parser.add_argument( @@ -221,12 +222,10 @@ def main(): logging.info(f"device: {device}") - sp = spm.SentencePieceProcessor() - sp.load(params.bpe_model) + lexicon = Lexicon(params.lang_dir) - # is defined in local/train_bpe_model.py - params.blank_id = sp.piece_to_id("") - params.vocab_size = sp.get_piece_size() + params.blank_id = lexicon.token_table[""] + params.vocab_size = max(lexicon.tokens) + 1 logging.info(params)