diff --git a/icefall/utils.py b/icefall/utils.py index 9dc718043..8ee62949d 100644 --- a/icefall/utils.py +++ b/icefall/utils.py @@ -1,5 +1,6 @@ -# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang -# Mingshuang Luo) +# Copyright 2021 Xiaomi Corp. (authors: Fangjun Kuang, +# Mingshuang Luo, +# Zengwei Yao) # # See ../../LICENSE for clarification regarding multiple authors # diff --git a/test/test_parse_timestamp.py b/test/test_parse_timestamp.py new file mode 100755 index 000000000..b3f06c54b --- /dev/null +++ b/test/test_parse_timestamp.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +# Copyright 2023 Xiaomi Corp. (authors: Zengwei Yao) +# +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from pathlib import Path + +import k2 +import sentencepiece as spm +import torch + +from icefall.lexicon import Lexicon +from icefall.utils import parse_bpe_timestamps_and_texts, parse_timestamps_and_texts + +ICEFALL_DIR = Path(__file__).resolve().parent.parent + + +def test_parse_bpe_timestamps_and_texts(): + lang_dir = ICEFALL_DIR / "egs/librispeech/ASR/data/lang_bpe_500" + if not lang_dir.is_dir(): + print(f"{lang_dir} does not exist.") + return + + sp = spm.SentencePieceProcessor() + sp.load(str(lang_dir / "bpe.model")) + + text = "HELLO WORLD" + token_ids = sp.encode(text, out_type=int) + # out_type=str: ['_HE', 'LL', 'O', '_WORLD'] + # out_type=int: [22, 58, 24, 425] + + # [22, 22, 58, 24, 0, 0, 425, 425, 425, 0, 0] + labels = ( + token_ids[0:1] * 2 + token_ids[1:3] + [0] * 2 + token_ids[3:4] * 3 + [0] * 2 + ) + # [22, 0, 58, 24, 0, 0, 425, 0, 0, 0, 0] + aux_labels = ( + token_ids[0:1] + + [0] + + token_ids[1:3] + + [0] * 2 + + token_ids[3:4] + + [0] * 4 + + [-1] + ) + + fsa = k2.linear_fsa(labels) + fsa.aux_labels = torch.tensor(aux_labels).to(torch.int32) + + fsa_vec = k2.create_fsa_vec([fsa]) + + utt_index_pairs, utt_words = parse_bpe_timestamps_and_texts(fsa_vec, sp) + assert utt_index_pairs[0] == [(0, 3), (6, 8)], utt_index_pairs[0] + assert utt_words[0] == ["HELLO", "WORLD"], utt_words[0] + + +def test_parse_timestamps_and_texts(): + lang_dir = ICEFALL_DIR / "egs/librispeech/ASR/data/lang_bpe_500" + if not lang_dir.is_dir(): + print(f"{lang_dir} does not exist.") + return + + lexicon = Lexicon(lang_dir) + + sp = spm.SentencePieceProcessor() + sp.load(str(lang_dir / "bpe.model")) + + text = "HELLO WORLD" + token_ids = sp.encode(text, out_type=int) + # out_type=str: ['_HE', 'LL', 'O', '_WORLD'] + # out_type=int: [22, 58, 24, 425] + + word_table = lexicon.word_table + word_ids = [word_table[s] for s in text.split()] # [79677, 196937] + + # [22, 22, 58, 24, 0, 0, 425, 425, 425, 0, 0] + labels = ( + token_ids[0:1] * 2 + token_ids[1:3] + [0] * 2 + token_ids[3:4] * 3 + [0] * 2 + ) + + # [[79677], [], [], [], [], [], [196937], [], [], [], [], []] + aux_labels = [word_ids[0:1]] + [[]] * 5 + [word_ids[1:2]] + [[]] * 5 + + fsa = k2.linear_fsa(labels) + fsa.aux_labels = k2.RaggedTensor(aux_labels) + + fsa_vec = k2.create_fsa_vec([fsa, fsa]) + + utt_index_pairs, utt_words = parse_timestamps_and_texts(fsa_vec, word_table) + assert utt_index_pairs[0] == [(0, 3), (6, 8)], utt_index_pairs[0] + assert utt_words[0] == ["HELLO", "WORLD"], utt_words[0] + + +if __name__ == "__main__": + test_parse_bpe_timestamps_and_texts() + test_parse_timestamps_and_texts()