diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 6397ca5..9004132 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -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): diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 77c77be..008c488 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -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', diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 4c3965e..0847ca6 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -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])