diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 787c5a3..635c9a6 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -68,7 +68,7 @@ def test_ulmfit_default_end_to_end(): def test_ulmfit_fastai_end_to_end(): """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset. """ - imdb, wt2 = get_test_data() + test_data, wt2 = get_test_data() lm_name = 'end-to-end-test-fastai' cuda_id = 0 exp = ulmfit.pretrain_lm.LMHyperParams( @@ -82,11 +82,13 @@ def test_ulmfit_fastai_end_to_end(): name=lm_name, ) exp.train_lm(num_epochs=1) + exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir) + exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, ) def test_ulmfit_fastai_bidir_end_to_end(): """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset. """ - imdb, wt2 = get_test_data() + test_data, wt2 = get_test_data() lm_name = 'end-to-end-test-fastai' cuda_id = 0 exp = ulmfit.pretrain_lm.LMHyperParams( @@ -101,11 +103,13 @@ def test_ulmfit_fastai_bidir_end_to_end(): name=lm_name, ) exp.train_lm(num_epochs=1) + exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir) + exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, ) def test_ulmfit_moses_fa_bidir_end_to_end(): """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset. """ - imdb, wt2 = get_test_data() + test_data, wt2 = get_test_data() lm_name = 'end-to-end-test-fastai' cuda_id = 0 exp = ulmfit.pretrain_lm.LMHyperParams( @@ -120,11 +124,13 @@ def test_ulmfit_moses_fa_bidir_end_to_end(): name=lm_name, ) exp.train_lm(num_epochs=1) + exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir) + exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, ) def test_ulmfit_sentencepiece_end_to_end(): """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset. """ - imdb, wt2 = get_test_data() + test_data, wt2 = get_test_data() lm_name = 'end-to-end-test-spm' cuda_id = 0 exp = ulmfit.pretrain_lm.LMHyperParams( @@ -138,10 +144,9 @@ def test_ulmfit_sentencepiece_end_to_end(): name=lm_name, ) exp.train_lm(num_epochs=1) - #assert exp.results['accuracy'] > 0.30 - - # NOTE: ds_pct is not available for sentencepiece -- tests are on the complete dataset - # sentencepiece for finetuning/classification is currently not implemented + # not supported yet + # exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir) + # exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, ) if __name__ == "__main__": diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 6169edf..e68ed95 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -219,6 +219,7 @@ class LMHyperParams: except FileNotFoundError: print("Running tokenization") + # wikitext is pretokenized with Moses pretokenized = Tokenizer(tok_func=BaseTokenizer, lang='en', pre_rules=None, post_rules=None) data_lm = TextLMDataBunch.from_df(path=self.cache_dir, train_df=read_file(trn_path), valid_df=read_file(val_path), tokenizer=pretokenized, @@ -237,7 +238,7 @@ class LMHyperParams: raise ValueError(f"self.tokenizer has wrong value {self.tokenizer}, Allowed values are taken from {Tokenizers}") itos, stoi, trn_path = data_lm.vocab.itos, data_lm.vocab.stoi, data_lm.path print('Size of vocabulary:', len(itos)) - print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)])) + print('First 20 words in vocab:', data_lm.vocab.itos[:20]) return data_lm @classmethod diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py index c327f3e..0782995 100644 --- a/ulmfit/train_clas.py +++ b/ulmfit/train_clas.py @@ -2,27 +2,46 @@ Train a classifier on top of a language model trained with `pretrain_lm.py`. Optionally fine-tune LM before. """ +from sacremoses import MosesTokenizer + import fastai import numpy as np import pickle +from fastai import * +from fastai.text import * + import torch from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner from fastai import fit_one_cycle, accuracy from fastai_contrib.data import LanguageModelType from fastai_contrib.learner import bilm_text_classifier_learner, bilm_learner, accuracy_fwd, accuracy_bwd -from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists +from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists, \ + get_sentencepiece from fastai.text.transform import Vocab + import fire from collections import Counter from pathlib import Path -from ulmfit.pretrain_lm import LMHyperParams +from ulmfit.pretrain_lm import LMHyperParams, Tokenizers +class MosesTokenizerFunc(BaseTokenizer): + "Wrapper around a MosesTokenizer to make it a `BaseTokenizer`." + def __init__(self, lang:str): + self.tok = MosesTokenizer(lang) + + def tokenizer(self, t:str) -> List[str]: + return self.tok.tokenize(t, return_str=False, escape=False) + + def add_special_cases(self, toks:Collection[str]): + for w in toks: + assert len(self.tokenizer(w))==1, f"Tokenizer is unable to keep {w} as one token!" class CLSHyperParams(LMHyperParams): # dir_path -> data/imdb/ + use_test_for_validation=False def __post_init__(self, *args, **kwargs): super().__post_init__(*args, **kwargs) @@ -66,11 +85,72 @@ class CLSHyperParams(LMHyperParams): learn = classifier_learner(data_clas, bptt=self.bptt, pad_token=PAD_TOKEN_ID, path=self.model_dir.parent, model_dir=self.model_dir.name, qrnn=self.qrnn, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, drop_mult=self.drop_mult) - - learn.metrics = [accuracy_fwd, accuracy_bwd] if self.bidir else [accuracy] + learn.true_wd = False + print("true_wd: ", learn.true_wd) return learn def load_cls_data(self, bs): + if self.dataset_dir.name == 'imdb': + return self.load_cls_data_imdb(bs) + else: + assert self.tokenizer is Tokenizers.MOSES, "XNLI does not support other tokenizers than Moses" + return self.load_cls_data_old_for_xnli(bs) + + def load_cls_data_imdb(self, bs): + trn_df = pd.read_csv(self.dataset_path / 'train.csv', header=None) + tst_df = pd.read_csv(self.dataset_path / 'test.csv', header=None) + + if self.use_test_for_validation: + val_len = max(int(len(tst_df) * 0.1), 2) + tst_len = len(tst_df) - val_len + val_df = trn_df[tst_len:] + else: + val_len = max(int(len(trn_df) * 0.1), 2) + trn_len = len(trn_df) - val_len + trn_df, val_df = trn_df[trn_len:], trn_df[trn_len:] + + if self.tokenizer is Tokenizers.SUBWORD: + #TODO Fix me to make sure it trains correct dictionary + args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv', self.name, vocab_size=self.max_vocab) + elif self.tokenizer is Tokenizers.MOSES: + args = dict(tokenizer=Tokenizer(tok_func=MosesTokenizerFunc, lang='en', pre_rules=[], post_rules=[])) + elif self.tokenizer is Tokenizers.MOSES_FA: + args = dict(tokenizer=Tokenizer(tok_func=MosesTokenizerFunc, lang='en')) # use default pre/post rules + elif self.tokenizer is Tokenizers.FASTAI: + args = dict() + else: + raise ValueError( + f"self.tokenizer has wrong value {self.tokenizer}, Allowed values are taken from {Tokenizers}") + + try: + data_lm = TextLMDataBunch.load(self.cache_dir, 'lm', lm_type=self.lm_type) + print("Tokenized data loaded") + except FileNotFoundError: + print("Running tokenization") + + # wikitext is pretokenized with Moses + + data_lm = TextLMDataBunch.from_df(path=self.cache_dir, train_df=trn_df, + valid_df=val_df, test_df=tst_df, + lm_type=self.lm_type, **args) + data_lm.save('lm') + + args['vocab'] = data_lm.vocab # make sure we use the same vocab for classifcation + try: + data_cls = TextClasDataBunch.load(self.cache_dir, '.', lm_type=self.lm_type) + print("Tokenized data loaded") + except FileNotFoundError: + print("Running tokenization") + data_cls = TextClasDataBunch.from_df(path=self.cache_dir, train_df=trn_df, + valid_df=val_df, test_df=tst_df, + **args) + data_cls.save('.') + print('Size of vocabulary:', len(data_lm.vocab.itos)) + print('First 20 words in vocab:', data_lm.vocab.itos[:20]) + return data_cls, data_lm + + + def load_cls_data_old_for_xnli(self, bs): tmp_dir = self.cache_dir tmp_dir.mkdir(exist_ok=True) vocab_file = tmp_dir / f'vocab_{self.lang}.pkl'