Merge pull request #8 from n-waves/sentencepiece

[WIP] Sub-word tokenization with sentencepiece
This commit is contained in:
Piotr Czapla
2018-11-16 23:35:01 +01:00
committed by GitHub
3 changed files with 103 additions and 25 deletions
+71 -2
View File
@@ -9,6 +9,12 @@ import torch
from tqdm import tqdm
import re
import csv
from functools import reduce
from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab, default_rules
from fastai.torch_core import *
import shutil
import pathlib
import tarfile
from sklearn import model_selection
@@ -28,10 +34,74 @@ XNLI_PATHS = {
TST: 'XNLI-1.0/xnli.test.tsv'
}
CLASSES = ['neg', 'pos', 'unsup']
number_match_re = re.compile(r'^([0-9]+[,.]?)+$')
number_split_re = re.compile(r'([,.])')
CLASSES = ['neg', 'pos', 'unsup']
class SentencepieceTokenizer(BaseTokenizer):
def __init__(self, model_dir:PathOrStr):
try:
import sentencepiece as spm
except ImportError:
raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')
self.tok = spm.SentencePieceProcessor()
self.tok.Load(str(pathlib.Path(model_dir) / 'spm.model'))
def tokenizer(self, t:str) -> List[str]:
return self.tok.EncodeAsPieces(t)
def add_special_cases(self, toks:Collection[str]):
pass
def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=None,
vocab_size:int=30000, model_type:str='unigram', input_sentence_size:int=1E7,
pad_idx:int=PAD_TOKEN_ID):
try:
import sentencepiece as spm
except ImportError:
raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')
path = pathlib.Path(path)
os.makedirs(path / 'models', exist_ok=True)
rules = rules if rules else default_rules
cache_name = 'tmp'
# load the text frmo the train tokens file
text = [line.rstrip('\n') for line in open(trn_path)]
text = list(filter(None, text))
if not os.path.isfile(path / 'models' / 'spm.model') or not os.path.isfile(path / f'itos_{name}.pkl'):
raw_text = reduce(lambda t, rule: rule(t), rules, '\n'.join(text))
raw_text_path = path / cache_name / 'all_text.txt'
with open(raw_text_path, 'w') as f:
f.write(raw_text)
sp_params = f'--input={raw_text_path} --pad_id={pad_idx} --unk_id=0' \
f'--character_coverage=1.0 --bos_id=-1 --eos_id=-1 ' \
f'--input_sentence_size={int(input_sentence_size)} ' \
f"--model_prefix={path / 'models' / 'spm'} " \
f'--vocab_size={vocab_size} --model_type={model_type} '
spm.SentencePieceTrainer.Train(sp_params)
with open(path / 'models' / 'spm.vocab', 'r') as f:
vocab = [line.split('\t')[0] for line in f.readlines()]
vocab[0] = UNK
vocab[pad_idx] = PAD
pickle.dump(vocab, open(path / 'models'/ f'itos_{name}.pkl', 'wb'))
vocab = Vocab(pickle.load(open(path / 'models'/ f'itos_{name}.pkl', 'rb')))
spt = SentencepieceTokenizer(path)
tokenizer = Tokenizer(tok_func=lambda lang: spt, rules=rules)
clear_cache_directory(path, cache_name)
return {'tokenizer': tokenizer, 'vocab': vocab}
def clear_cache_directory(path:PathOrStr, cache_name:str='tmp'):
path = pathlib.Path(path)
shutil.rmtree(path / cache_name)
def get_texts(path):
@@ -225,7 +295,6 @@ def read_clas_data(dir_path, dataset, lang) -> Tuple[Dict[str, List[List[str]]],
toks[VAL], lbls[VAL] = processor(dir_path, lang, VAL)
return toks, lbls
def replace_number(token):
"""Replaces a number and returns a list of one or multiple tokens."""
if number_match_re.match(token):
+4 -1
View File
@@ -23,11 +23,13 @@ def check_data_exists():
def test_pretrain_lm():
imdb,wt2 = check_data_exists()
lm_name="end-to-end-test-quick"
cuda_id=0
results = ulmfit.pretrain_lm.pretrain_lm(
dir_path=wt2,
lang='en',
cuda_id=cuda_id,
qrnn=True,
clean=True,
subword=False,
max_vocab=1000,
bs=80,
num_epochs=1,
@@ -40,6 +42,7 @@ def test_pretrain_lm():
data_dir=get_data_folder(),
lang='en', pretrain_name=lm_name, model_dir=wt2/'models',
qrnn=True,
cuda_id=cuda_id,
fine_tune=True,
max_vocab=1000,
bs=20, bptt=70, name=lm_name+'-imdb-clas',
+28 -22
View File
@@ -12,7 +12,7 @@ from fastai import *
from fastai.text import *
import torch
from fastai_contrib.utils import read_file, read_whitespace_file,\
DataStump, validate, PAD, UNK
DataStump, validate, PAD, UNK, get_sentencepiece
import pickle
@@ -24,16 +24,15 @@ from collections import Counter
# conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92]
# cupy needs to be installed for QRNN
def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vocab=60000,
bs=70, bptt=70, name='wt-103', num_epochs=10, ds_pct=1.0):
"""
:param dir_path: The path to the directory that contains wiki text
:param dir_path: The path to the directory of the file.
:param lang: the language unicode
:param cuda_id: The id of the GPU. Uses GPU 0 by default or no GPU when
run on CPU.
:param qrrn: Use a QRNN. Requires installing cupy.
:param clean: Train on the clean
:param qrnn: Use a QRNN. Requires installing cupy.
:param subword: Use sub-word tokenization on the cleaned data.
:param max_vocab: The maximum size of the vocabulary.
:param bs: The batch size.
:param bptt: The back-propagation-through-time sequence length.
@@ -63,7 +62,20 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
for path_ in [trn_path, val_path, tst_path]:
assert path_.exists(), f'Error: {path_} does not exist.'
if clean:
if subword:
# apply sentencepiece tokenization
trn_path = dir_path / f'{lang}.wiki.train.tokens'
val_path = dir_path / f'{lang}.wiki.valid.tokens'
read_file(trn_path, 'train')
read_file(val_path, 'valid')
sp = get_sentencepiece(dir_path, trn_path, name)
data_lm = TextLMDataBunch.from_csv(dir_path, **sp)
itos = data_lm.train_ds.vocab.itos
stoi = data_lm.train_ds.vocab.stoi
else:
# read the already whitespace separated data without any preprocessing
trn_tok = read_whitespace_file(trn_path)
val_tok = read_whitespace_file(val_path)
@@ -81,20 +93,21 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
vocab = Vocab(itos)
stoi = vocab.stoi
# save vocabulary
print(f"Saving vocabulary as {dir_path / model_dir}")
results['itos_fname'] = dir_path / model_dir / f'itos_{name}.pkl'
with open(results['itos_fname'], 'wb') as f:
pickle.dump(itos, f)
trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok])
val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_tok])
# data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos))
data_lm = TextLMDataBunch.from_ids(path=dir_path, vocab=vocab, train_ids=trn_ids,
valid_ids=val_ids, bs=bs, bptt=bptt)
else:
# apply fastai preprocessing and tokenization
read_file(trn_path, 'train')
read_file(val_path, 'valid')
data_lm = TextLMDataBunch.from_csv(dir_path, max_vocab=max_vocab)
itos = data_lm.train_ds.vocab.itos
stoi = data_lm.train_ds.vocab.stoi
print('Size of vocabulary:', len(itos))
print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)]))
@@ -120,20 +133,13 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99))
learn.true_wd = False
# save vocabulary
print(f"Saving vocabulary as {dir_path / model_dir}")
results['itos_fname'] = dir_path / model_dir / f'itos_{name}.pkl'
with open(results['itos_fname'], 'wb') as f:
pickle.dump(itos, f)
fit_one_cycle(learn, num_epochs, 5e-3, (0.8, 0.7), wd=1e-7)
if clean and max_vocab is None:
if not subword and max_vocab is None:
# only if we use the unpreprocessed version and the full vocabulary
# are the perplexity results comparable to previous work
print(f"Validating model performance with test tokens from: {trn_path}")
tst_tok = read_whitespace_file(trn_path)
tst_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in tst_tok])