mirrors/icefall
1
0
Fork 0
mirror of https://github.com/k2-fsa/icefall.git synced 2025-09-06 23:54:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

minor fixes

Browse Source
This commit is contained in:
jinzr 2023-11-23 17:41:14 +08:00
parent a7b4e8f52f
commit e339dbd3d3
2 changed files with 16 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
egs/multi_zh_en/ASR/zipformer/decode.py
View File

@ -130,7 +130,6 @@ from icefall.checkpoint import (
from icefall.lexicon import Lexicon from icefall.lexicon import Lexicon
from icefall.utils import ( from icefall.utils import (
AttributeDict, AttributeDict,
make_pad_mask,
setup_logger, setup_logger,
store_transcripts, store_transcripts,
str2bool, str2bool,

35
egs/multi_zh_en/ASR/zipformer/pretrained.py
View File

@ -115,6 +115,7 @@ from typing import List
import k2 import k2
import kaldifeat import kaldifeat
import sentencepiece as spm
import torch import torch
import torchaudio import torchaudio
from beam_search import ( from beam_search import (
@ -126,7 +127,7 @@ from export import num_tokens
from torch.nn.utils.rnn import pad_sequence from torch.nn.utils.rnn import pad_sequence
from train import add_model_arguments, get_model, get_params from train import add_model_arguments, get_model, get_params
from icefall.utils import make_pad_mask from icefall import smart_byte_decode
def get_parser(): def get_parser():
@ -144,9 +145,9 @@ def get_parser():
) )
parser.add_argument( parser.add_argument(
"--tokens", "--bpe-model",
type=str, type=str,
help="""Path to tokens.txt.""", help="""Path to byte-level bpe model.""",
) )
parser.add_argument( parser.add_argument(
@ -263,11 +264,13 @@ def main():
params.update(vars(args)) params.update(vars(args))
token_table = k2.SymbolTable.from_file(params.tokens) sp = spm.SentencePieceProcessor()
sp.load(params.bpe_model)
params.blank_id = token_table["<blk>"] # <blk> and <unk> are defined in local/train_bpe_model.py
params.unk_id = token_table["<unk>"] params.blank_id = sp.piece_to_id("<blk>")
params.vocab_size = num_tokens(token_table) + 1 params.unk_id = sp.piece_to_id("<unk>")
params.vocab_size = sp.get_piece_size()
logging.info(f"{params}") logging.info(f"{params}")
@ -326,12 +329,6 @@ def main():
msg = f"Using {params.method}" msg = f"Using {params.method}"
logging.info(msg) logging.info(msg)
def token_ids_to_words(token_ids: List[int]) -> str:
text = ""
for i in token_ids:
text += token_table[i]
return text.replace("▁", " ").strip()
if params.method == "fast_beam_search": if params.method == "fast_beam_search":
decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device) decoding_graph = k2.trivial_graph(params.vocab_size - 1, device=device)
hyp_tokens = fast_beam_search_one_best( hyp_tokens = fast_beam_search_one_best(
@ -343,8 +340,8 @@ def main():
max_contexts=params.max_contexts, max_contexts=params.max_contexts,
max_states=params.max_states, max_states=params.max_states,
) )
for hyp in hyp_tokens: for hyp in sp.decode(hyp_tokens):
hyps.append(token_ids_to_words(hyp)) hyps.append(smart_byte_decode(hyp).split())
elif params.method == "modified_beam_search": elif params.method == "modified_beam_search":
hyp_tokens = modified_beam_search( hyp_tokens = modified_beam_search(
model=model, model=model,
@ -353,16 +350,16 @@ def main():
beam=params.beam_size, beam=params.beam_size,
) )
for hyp in hyp_tokens: for hyp in sp.decode(hyp_tokens):
hyps.append(token_ids_to_words(hyp)) hyps.append(smart_byte_decode(hyp).split())
elif params.method == "greedy_search" and params.max_sym_per_frame == 1: elif params.method == "greedy_search" and params.max_sym_per_frame == 1:
hyp_tokens = greedy_search_batch( hyp_tokens = greedy_search_batch(
model=model, model=model,
encoder_out=encoder_out, encoder_out=encoder_out,
encoder_out_lens=encoder_out_lens, encoder_out_lens=encoder_out_lens,
) )
for hyp in hyp_tokens: for hyp in sp.decode(hyp_tokens):
hyps.append(token_ids_to_words(hyp)) hyps.append(smart_byte_decode(hyp).split())
else: else:
raise ValueError(f"Unsupported method: {params.method}") raise ValueError(f"Unsupported method: {params.method}")