mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-27 10:44:19 +00:00
119 lines
3.5 KiB
Python
Executable File
119 lines
3.5 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
# Copyright (c) 2021 Xiaomi Corporation (authors: Daniel Povey)
|
|
|
|
"""
|
|
|
|
This script takes a `bpe.model` and a text file such as
|
|
`download/lm/librispeech-lm-norm.txt`,
|
|
and outputs the LM training data to a supplied directory such
|
|
as data/lm_training_data_bpe_5000. The format is as follows:
|
|
|
|
It creates a PyTorch archive (.pt file), say data/lm_training.pt, which is a
|
|
representation of a dict with the following format:
|
|
|
|
'words' -> a k2.RaggedTensor of two axes [word][token] with dtype torch.int32
|
|
containing the BPE representations of each word, indexed by
|
|
integer word ID. (These integer word IDS are present in
|
|
'lm_data'). The sentencepiece object can be used to turn the
|
|
words and BPE units into string form.
|
|
'data' -> a k2.RaggedTensor of two axes [sentence][word] with dtype
|
|
torch.int32 containing all the sentences, as word-ids (we don't
|
|
output the string form of this directly but it can be worked out
|
|
together with 'words' and the bpe.model).
|
|
|
|
"""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
from typing import Dict, List, Tuple
|
|
|
|
import k2
|
|
import sentencepiece as spm
|
|
import torch
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"bpe_model",
|
|
type=str,
|
|
help="""Input BPE model, e.g. data/lang_bpe/bpe.model""",
|
|
)
|
|
parser.add_argument(
|
|
"lm_data",
|
|
type=str,
|
|
help="""Input LM training data as text, e.g. data/downloads/lm/librispeech-lm-norm.txt""",
|
|
)
|
|
parser.add_argument(
|
|
"lm_archive",
|
|
type=str,
|
|
help="""Path to output archive, e.g. lm_data.pt; look at the source of this script to see the format.""",
|
|
)
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def main():
|
|
args = get_args()
|
|
|
|
sp = spm.SentencePieceProcessor()
|
|
sp.load(args.bpe_model)
|
|
|
|
# word2index is a dictionary from words to integer ids. No need to reserve
|
|
# space for epsilon, etc.; the words are just used as a convenient way to
|
|
# compress the sequences of BPE pieces.
|
|
word2index = dict()
|
|
|
|
words2bpe = [] # Will be a list-of-list-of-int, representing BPE pieces.
|
|
|
|
sentences = [] # Wil be a list-of-list-of-int, representing word-ids.
|
|
|
|
with open(args.lm_data) as f:
|
|
while True:
|
|
line = f.readline()
|
|
if line == "":
|
|
break
|
|
line_words = line.split()
|
|
for w in line_words:
|
|
if not w in word2index:
|
|
w_bpe = sp.Encode(w)
|
|
word2index[w] = len(words2bpe)
|
|
words2bpe.append(w_bpe)
|
|
sentences.append([word2index[w] for w in line_words])
|
|
|
|
output = dict()
|
|
output["words"] = k2.ragged.RaggedTensor(words2bpe)
|
|
output["data"] = k2.ragged.RaggedTensor(sentences)
|
|
|
|
torch.save(output, args.lm_archive)
|
|
print(f"Saved to {args.lm_archive}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|
|
|
|
# This was tested as follows.
|
|
# cat > foo <<EOF
|
|
# THING TWO
|
|
# ZOOLOGY
|
|
# EOF
|
|
#
|
|
# local/prepare_lm_training_data.py data/lang_bpe/bpe.model foo bar.pt
|
|
#
|
|
# python3
|
|
# Python 3.8.0 (default, Oct 28 2019, 16:14:01)
|
|
# [GCC 8.3.0] on linux
|
|
# Type "help", "copyright", "credits" or "license" for more information.
|
|
# >>> import k2
|
|
# >>> import sentencepiece as spm
|
|
# >>> sp = spm.SentencePieceProcessor()
|
|
# >>> sp.load('data/lang_bpe/bpe.model')
|
|
# True
|
|
# >>> import torch
|
|
# >>> d = torch.load('bar.pt')
|
|
# >>> sp.Decode(k2.ragged.to_list(k2.index(d['words'], d['data'])))
|
|
# ['THING TWO', 'ZOOLOGY']
|
|
# >>>
|