fix export for ncnn

This commit is contained in:
marcoyang 2023-02-13 12:04:29 +08:00
parent 72f5c5e17f
commit 4d4a2c08f7

View File

@ -60,6 +60,7 @@ from icefall.checkpoint import (
find_checkpoints, find_checkpoints,
load_checkpoint, load_checkpoint,
) )
from icefall.lexicon import Lexicon
from icefall.utils import setup_logger, str2bool from icefall.utils import setup_logger, str2bool
@ -106,10 +107,10 @@ def get_parser():
) )
parser.add_argument( parser.add_argument(
"--bpe-model", "--lang-dir",
type=str, type=str,
default="data/lang_bpe_500/bpe.model", default="data/lang_char",
help="Path to the BPE model", help="Path to the lang",
) )
parser.add_argument( parser.add_argument(
@ -221,12 +222,10 @@ def main():
logging.info(f"device: {device}") logging.info(f"device: {device}")
sp = spm.SentencePieceProcessor() lexicon = Lexicon(params.lang_dir)
sp.load(params.bpe_model)
# <blk> is defined in local/train_bpe_model.py params.blank_id = lexicon.token_table["<blk>"]
params.blank_id = sp.piece_to_id("<blk>") params.vocab_size = max(lexicon.tokens) + 1
params.vocab_size = sp.get_piece_size()
logging.info(params) logging.info(params)