[WIP] Sub-word tokenization with sentencepiece

For resolving issue #3.

@eisenjulian let's use this branch. I've written some skeleton code (untested currently) to be used around your tokenizer. You can commit that into fastai_contrib for our purpose.
This commit is contained in:
Aayush
2018-11-15 23:01:49 +05:30
committed by GitHub
parent 898c9255e0
commit 3e8f9f8b5f
+97 -5
View File
@@ -11,6 +11,7 @@ import json
from shutil import copyfile
from sacremoses import MosesTokenizer
from fastai_contrib.utils import get_sentencepiece, replace_number, UNK
def get_texts(root):
@@ -27,10 +28,12 @@ def get_texts(root):
yield text
def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'):
def write_wikitext(file_path, text_iter, tok, num_tokens, mode='w'):
total_num_tokens = 0
print(f'Writing to {file_path}...')
i = 0
with open(file_path, mode, encoding='utf-8') as f_out:
for i, text in enumerate(text_iter):
@@ -39,7 +42,7 @@ def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'):
paragraphs = text.split('\n')
for paragraph in paragraphs:
tokenized = mt.tokenize(paragraph.strip(), return_str=True)
tokenized = tok.tokenize(paragraph.strip(), return_str=True)
tokenized_paragraphs.append(tokenized)
tokens = tokenized.split(' ') # split on whitespace to keep newlines
@@ -64,6 +67,70 @@ def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'):
print('{}. # documents: {:,}. # tokens: {:,}.'.format(
file_path, i, total_num_tokens))
def build_vocab(file_path, cutoff=3):
counter = Counter()
with open(file_path, 'r', encoding='utf-8') as f:
for i, line in enumerate(f):
tokens = line.strip().split(' ') + ['<eos>']
counter.update(tokens)
vocab = {}
in_vocab_count = 0
OOV_count = 0
for token, count in counter.most_common():
if count >= cutoff:
vocab[token] = count
in_vocab_count += count
else:
OOV_count += count
print('OOV ratio: %.4f.' % (OOV_count / (in_vocab_count + OOV_count)))
return vocab
def limit_vocab(unk_path, vocab):
"""
https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-limit_vocab-py
:param unk_path:
:param vocab:
:return:
"""
temp_file_path = unk_path.with_name(unk_path.name + '.temp')
total_num_tokens = 0
print(f'Limiting vocab in {unk_path}. Writing to {unk_path}.')
with open(unk_path, 'r', encoding='utf-8') as f_in, open(temp_file_path, 'w', encoding='utf-8') as f_out:
for line in f_in:
tokens = [x for x in line.strip().split(' ') if x]
tokens = [token if token in vocab else UNK for token in tokens]
# Ensures there's a space between tokens, including the last word,
# newline, and the first word of the next line
tokens = tokens + ['\n']
total_num_tokens += len(tokens)
tokens = [''] + tokens
line = ' '.join(tokens)
f_out.write(line)
print(f'{unk_path.name}. # of tokens: {total_num_tokens}')
temp_file_path.replace(unk_path)
def replace_numbers(text_iter, unk_path):
"""
Replace numbers as in Smerity's script:
https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-post_process-py
:param file_path:
:param unk_path:
:return:
"""
print(f'Replacing numbers in file. Writing to {unk_path}.')
with open(unk_path, 'w', encoding='utf-8') as f:
for text in text_iter:
raw_tokens = line.strip().split(' ')
tokens = []
for token in raw_tokens:
tokens.append(replace_number(token))
# Starting each line with a blank line is required
# Some systems replace \n with <eos> and assume, like in PTB, everything is space separated
tokens = [''] + tokens + ['\n']
line = ' '.join(tokens)
f.write(line)
def main(args):
@@ -72,7 +139,13 @@ def main(args):
assert input_path.exists(), f'Error: {input_path} does not exist.'
output.mkdir(exist_ok=True)
mt = MosesTokenizer(args.lang)
if args.subword:
# TO DO load the text corpus
# TO DO make get_sentencepiece return path to spm model
spm_path = get_sentencepiece(output, corpus)
tok = SentencepieceTokenizer(spm_path)
else:
tok = MosesTokenizer(args.lang)
sml_wiki = output / f'{args.lang}-2'
lrg_wiki = output / f'{args.lang}-100'
@@ -84,17 +157,33 @@ def main(args):
splits = ['train', 'valid', 'test']
token_nums = [2000000, 200000, 200000]
for split, token_num in zip(splits, token_nums):
# TO DO maybe replace the numbers before tokenizing
unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk'
replace_numbers(text_iter, unk_path)
sml_file_path = sml_wiki / f'{args.lang}.wiki.{split}.tokens'
write_wikitext(sml_file_path, text_iter, mt, token_num)
write_wikitext(sml_file_path, text_iter, tok, token_num)
lrg_file_path = lrg_wiki / f'{args.lang}.wiki.{split}.tokens'
# copy the content of the small file to the large file
print(f'Copying {sml_file_path} to {lrg_file_path}.')
copyfile(sml_file_path, lrg_file_path)
sml_vocab = build_vocab(sml_wiki_train)
print(f'{args.lang}-2 vocab size: {len(sml_vocab)}')
lrg_vocab = build_vocab(lrg_wiki_train)
print(f'{args.lang}-100 vocab size: {len(lrg_vocab)}')
# add the new articles to the existing ones
lrg_wiki_train = lrg_wiki / f'{args.lang}.wiki.train.tokens'
write_wikitext(lrg_wiki_train, text_iter, mt, 98000000, mode='a')
write_wikitext(lrg_wiki_train, text_iter, tok, 98000000, mode='a')
# replace words not in the vocab with <unk>
if not args.subword:
for wiki, vocab in zip([sml_wiki, lrg_wiki], [sml_vocab, lrg_vocab]):
for split in splits:
unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk'
limit_vocab(unk_path, vocab)
if __name__ == '__main__':
@@ -110,5 +199,8 @@ if __name__ == '__main__':
parser.add_argument('-l', '--lang', required=True,
help='the iso code of the language of the Wikipedia '
'documents, e.g. en, fr, de, etc.')
parser.add_argument('-sw', '--subword', default=False,
help='set to use sub-word tokenization (sentencepiece)'
'default tokenization method is Moses.')
args = parser.parse_args()
main(args)