From 3e8f9f8b5f057123918dfdef7eeb55f8320ccd7a Mon Sep 17 00:00:00 2001 From: Aayush Date: Thu, 15 Nov 2018 23:01:49 +0530 Subject: [PATCH] [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. --- ulmfit/create_wikitext.py | 102 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 97 insertions(+), 5 deletions(-) diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py index 5d24284..b6f3dd9 100644 --- a/ulmfit/create_wikitext.py +++ b/ulmfit/create_wikitext.py @@ -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(' ') + [''] + 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 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 + 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)