mirror of
https://github.com/wassname/multifit.git
synced 2026-08-22 12:10:15 +08:00
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""
|
|
Script to post-process WikiText files created with `create_wikitext.py`.
|
|
Creates additional files where words not in the training data are replaced
|
|
with <UNK> and numbers are modified with a regex.
|
|
"""
|
|
import argparse
|
|
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
import fire
|
|
|
|
from fastai_contrib.utils import replace_number, UNK
|
|
|
|
|
|
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(file_path, 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_path}. Writing to {unk_path}.')
|
|
with open(file_path, 'r', encoding='utf-8') as f_in, open(unk_path, 'w', encoding='utf-8') as f_out:
|
|
for line in f_in:
|
|
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_out.write(line)
|
|
|
|
|
|
def postprocess_wikitext(path, lang):
|
|
|
|
wiki_path = Path(path)
|
|
assert wiki_path.exists(), f'Error: {wiki_path} does not exist.'
|
|
dest_path = wiki_path.parent / (wiki_path.name + "-unk")
|
|
dest_path.mkdir(exist_ok=True)
|
|
splits = ['train', 'valid', 'test']
|
|
for split in splits:
|
|
# replace numbers with placeholders
|
|
file_path = wiki_path / f'{lang}.wiki.{split}.tokens'
|
|
assert file_path.exists(), f"Error: {file_path} does not exist."
|
|
unk_path = dest_path / file_path.name
|
|
replace_numbers(file_path, unk_path)
|
|
|
|
# replace words not in the vocab with <unk>
|
|
wiki_train = dest_path / f'{lang}.wiki.train.tokens'
|
|
vocab = build_vocab(wiki_train)
|
|
print(f'{wiki_path} vocab size: {len(vocab)}')
|
|
for split in splits:
|
|
unk_path = dest_path / f'{lang}.wiki.{split}.tokens'
|
|
limit_vocab(unk_path, vocab)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
fire.Fire(postprocess_wikitext)
|