mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-13 20:12:24 +00:00
119 lines
3.4 KiB
Python
Executable File
119 lines
3.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
# Copyright (c) 2021 Xiaomi Corporation (authors: Fangjun Kuang, 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._RaggedInt containing the BPE representations of each word, inexed 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._RaggedInt 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.create_ragged2(words2bpe)
|
|
output['data'] = k2.ragged.create_ragged2(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']
|
|
#>>>
|