mirror of
https://github.com/wassname/multifit.git
synced 2026-09-11 12:20:41 +08:00
[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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user