mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-09-05 15:14:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix export for wenetspeech

Browse Source
This commit is contained in:
Fangjun Kuang 2024-01-26 11:59:42 +08:00
parent c72bba1e0b
commit c65b734838
1 changed files with 8 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
egs/wenetspeech/ASR/pruned_transducer_stateless2/export.py
View File

@ -91,14 +91,14 @@ import argparse
import logging import logging
from pathlib import Path from pathlib import Path
import k2
import torch import torch
import torch.nn as nn import torch.nn as nn
from scaling_converter import convert_scaled_to_non_scaled from scaling_converter import convert_scaled_to_non_scaled
from train import get_params, get_transducer_model from train import get_params, get_transducer_model
from icefall.checkpoint import average_checkpoints, load_checkpoint from icefall.checkpoint import average_checkpoints, load_checkpoint
from icefall.lexicon import Lexicon from icefall.utils import num_tokens, str2bool
from icefall.utils import str2bool
def get_parser(): def get_parser():
@ -133,10 +133,10 @@ def get_parser():
) )
parser.add_argument( parser.add_argument(
"--lang-dir", "--tokens",
type=str, type=str,
default="data/lang_char", default="data/lang_char/tokens.txt",
help="The lang dir", help="Path to the tokens.txt",
) )
parser.add_argument( parser.add_argument(
@ -313,10 +313,9 @@ def main():
logging.info(f"device: {device}") logging.info(f"device: {device}")
lexicon = Lexicon(params.lang_dir) token_table = k2.SymbolTable.from_file(params.tokens)
params.blank_id = token_table["<blk>"]
params.blank_id = 0 params.vocab_size = num_tokens(token_table) + 1
params.vocab_size = max(lexicon.tokens) + 1
logging.info(params) logging.info(params)