From 70c74a1cc5f9338fbd05232d6ffec00253df9240 Mon Sep 17 00:00:00 2001 From: Piotr Czapla Date: Tue, 15 Oct 2019 04:17:54 +0200 Subject: [PATCH] Refactor tokenization So that we can instantiate tokenization before we know what dataset we want to use it on. Previously it was tidly copuled. --- README.md | 6 ++ sotabench/sotabench.py | 4 +- ulmfit/datasets/__init__.py | 2 +- ulmfit/datasets/dataset.py | 148 +++++++++++++++++++++--------------- ulmfit/training.py | 69 ++++++++++++----- 5 files changed, 147 insertions(+), 82 deletions(-) diff --git a/README.md b/README.md index 668a790..374dd07 100644 --- a/README.md +++ b/README.md @@ -12,3 +12,9 @@ You can evaulate any model with the following command: ```bash python -m ulmfit load data/mldoc/de-1/models/fsp15k/multfit_fp16 classifier validate data/mldoc/de-1 ``` + + +```bash +python -m ulmfit new multifit_fp16_nl3 pretrain-lm train- data/wiki/wikitext-103 + +``` \ No newline at end of file diff --git a/sotabench/sotabench.py b/sotabench/sotabench.py index f1adebe..166e831 100644 --- a/sotabench/sotabench.py +++ b/sotabench/sotabench.py @@ -25,9 +25,7 @@ def evaluate(pretrained_name): wikitext_folder = WikiText103Evaluator.dataset.get_path(local_root="unused") else: wikitext_folder = untar_data(URLs.WIKITEXT) - ds = model.arch.dataset(wikitext_folder) - - ds.use_base_model_subword_vocabulary(model.pretrain_lm.experiment_path) + ds = model.arch.dataset(wikitext_folder, tokenizer=model.pretrain_lm.tokenizer) test_df = ds.read_data(ds.tst_path) data_lm = ds.databunch_from_df(TextLMDataBunch, test_df, test_df, bs=20, bptt=70) diff --git a/ulmfit/datasets/__init__.py b/ulmfit/datasets/__init__.py index a0bcc0d..d5e918d 100644 --- a/ulmfit/datasets/__init__.py +++ b/ulmfit/datasets/__init__.py @@ -1 +1 @@ -from .dataset import Dataset, ULMFiTDataset, read_clas_csv, read_wiki_articles \ No newline at end of file +from .dataset import Dataset, ULMFiTDataset, read_clas_csv, read_wiki_articles, ULMFiTTokenizer \ No newline at end of file diff --git a/ulmfit/datasets/dataset.py b/ulmfit/datasets/dataset.py index c6cfad8..30a6c45 100644 --- a/ulmfit/datasets/dataset.py +++ b/ulmfit/datasets/dataset.py @@ -184,47 +184,15 @@ class Dataset: @dataclass class ULMFiTDataset(Dataset): - tokenizer: str = 'f' - max_vocab: int = 60000 + tokenizer: Tokenizer = None cache_path: Path = None def __post_init__(self): super().__post_init__() if self.cache_path is None: - tokenizer_prefix = f"{self.tokenizer}{self.max_vocab // 1000}k" - self.cache_path = self.dataset_path / "models" / tokenizer_prefix + self.cache_path = self.dataset_path / "models" / self.tokenizer.prefix self._vocab = None - def use_base_model_subword_vocabulary(self, base_lm_path: Path): - """ - In case of subwoard vocabularies reuse the base model vocabulary during tokenization. - For word tokenization we still generate new vocabulary for each dataset, - and we expect finetuning to handle the conversion - """ - def copy_sp(path): - print(f"Copy sp model from {path} to {self.cache_path}") - shutil.copy(str(path / 'itos.pkl'), str(self.cache_path)) - shutil.copy(str(path / 'spm.model'), str(self.cache_path)) - shutil.copy(str(path / 'spm.vocab'), str(self.cache_path)) - - # reuse base model sentencepiece vocabulary - self.cache_path.mkdir(exist_ok=True, parents=True) - if base_lm_path is None or base_lm_path.parent.resolve() == self.cache_path.resolve(): - return - - if (base_lm_path.parent / 'spm.vocab').exists(): - copy_sp(base_lm_path.parent) - - if (base_lm_path / 'spm.vocab').exists(): - copy_sp(base_lm_path) - - # TODO: implement / maybe put the vocabulary md5 to the file names and keep spm models together? - # sp12k/7599013a8ce538b2e3d4405684221ecaf26bcba1.lm - # sp12k/7599013a8ce538b2e3d4405684221ecaf26bcba1-vocab.link - # then we don't need the use_vocabulary, we can just use load_lm_databunch(using_vacab=XXX) - # we could put spm.model and spm.vocab to the model folder it self then, and copy /link it when we use the orignal model - # for the time being we can simply compy the spm.model on the right spot and raise an error ir the two are different? - def load_lm_databunch(self, bs, bptt): lm_suffix = bptt if bptt != 70 else "" lm_suffix += self.use_tst_for_lm if "" else "-notst" @@ -278,18 +246,56 @@ class ULMFiTDataset(Dataset): return databunch def databunch_from_df(self, bunch_class, train_df, valid_df, **args): - args.update(**self.get_processor(ds_need_moses=not self.uses_moses)) # TODO depends on the previous model + args.update(**self.tokenizer.get_fastai_config(dataset_uses_moses=self.uses_moses)) # TODO depends on the previous model databunch = make_data_bunch_from_df(cls=bunch_class, path=self.cache_path, train_df=train_df, valid_df=valid_df, - max_vocab=self.max_vocab, mark_fields=True, text_cols=list(train_df.columns.values)[1:], **args) return databunch - def get_processor(self, ds_need_moses, add_open_file_processor=False): + +@dataclass +class ULMFiTTokenizer: + arch: Any # should be ULMFiTArchitecture, we use Any to avoid circular dependencies between imports + pretrained_path: Path = None + + def __post_init__(self): + self.temp_dir = None + if self.pretrained_path is None: + self.temp_dir = tempfile.TemporaryDirectory() + self.pretrained_path = Path(self.temp_dir.name) + + def save(self, new_path: Path, learn: Learner): + """ + In case of subwoard vocabularies reuse the base model vocabulary during tokenization. + For word tokenization we still generate new vocabulary for each dataset, + and we expect finetuning to handle the conversion + """ + + def copy_sp(path): + print(f"Copy sp model from {path} to {new_path}") + shutil.copy(str(path / 'itos.pkl'), str(new_path)) + shutil.copy(str(path / 'spm.model'), str(new_path)) + shutil.copy(str(path / 'spm.vocab'), str(new_path)) + + # reuse base model sentencepiece vocabulary + new_path.mkdir(exist_ok=True, parents=True) + if self.pretrained_path is None or self.pretrained_path.parent.resolve() == new_path.resolve(): + return + + if (self.pretrained_path.parent / 'spm.vocab').exists(): + copy_sp(self.pretrained_path.parent) + + if (self.pretrained_path / 'spm.vocab').exists(): + copy_sp(self.pretrained_path) + + with (new_path / "itos.pkl").open('wb') as f: + pickle.dump(learn.data.vocab.itos, f) + + def get_fastai_config(self, dataset_uses_moses=False, add_open_file_processor=False): return { 'fsp': self._get_processor_sentence_piece, 'f': self._get_processor_pure_fastai, @@ -299,44 +305,64 @@ class ULMFiTDataset(Dataset): 'sp': self._get_processor_sentence_piece, # deprecated 'v': self._get_processor_pure_moses, # deprecated 'vf': self._get_processor_moses_fastai, # deprecated - }.get(self.tokenizer)(ds_need_moses, add_open_file_processor) + }.get(self.arch.tokenizer)(dataset_uses_moses, add_open_file_processor) - def _get_processor_sentence_piece(self, ds_need_moses, add_open_file_processor=False): - moses_preproc = [MosesPreprocessingFunc(self.lang)] if ds_need_moses else [] + @property + def prefix(self): + return f"{self.arch.tokenizer}{self.arch.max_vocab // 1000}k" - sp_model = self.cache_path / 'spm.model' + def _get_processor_sentence_piece(self, ds_uses_moses, add_open_file_processor=False): + moses_preproc = [MosesPreprocessingFunc(self.arch.lang)] if not ds_uses_moses else [] + + sp_model = self.pretrained_path / 'spm.model' if not sp_model.is_file(): sp_model = None - sp_vocab = self.cache_path / 'spm.vocab' + sp_vocab = self.pretrained_path / 'spm.vocab' if not sp_vocab.is_file(): sp_vocab = None processor = SPProcessor2( pre_rules=moses_preproc + defaults.text_pre_rules, mark_fields=True, - vocab_sz=self.max_vocab, + vocab_sz=self.arch.max_vocab, sp_model=sp_model, sp_vocab=sp_vocab, - lang=self.lang, - tmp_dir=self.cache_path.absolute() # absolute make sure that dataset path is not added as prefix + lang=self.arch.lang, + tmp_dir=self.pretrained_path.absolute() # absolute make sure that dataset path is not added as prefix ) openfile = [OpenFileProcessor()] if add_open_file_processor else [] return {'processor': openfile + [ processor ]} - def _get_processor_pure_moses(self, ds_need_moses, add_open_file_processor=False): - moses_preproc = [MosesPreprocessingFunc(self.lang)] if ds_need_moses else [] - return dict(tokenizer=Tokenizer(tok_func=BaseTokenizer, - lang=self.lang, - pre_rules=moses_preproc, - post_rules=[])) + def _default_processor(self, fastai_tokenizer): + fastai_tokenizer = Tokenizer(SpacyTokenizer, self.arch.lang) + return [TokenizeProcessor(tokenizer=fastai_tokenizer), NumericalizeProcessor(max_vocab=self.arch.max_vocab)] - def _get_processor_moses_fastai(self, ds_need_moses, add_open_file_processor=False): - moses_preproc = [MosesPreprocessingFunc(self.lang)] if ds_need_moses else [] - return dict(tokenizer=Tokenizer(tok_func=BaseTokenizer, - lang=self.lang, - pre_rules=moses_preproc + defaults.text_pre_rules, - post_rules=defaults.text_post_rules)) + def _get_processor_pure_moses(self, ds_uses_moses, add_open_file_processor=False): + moses_preproc = [MosesPreprocessingFunc(self.arch.lang)] if not ds_uses_moses else [] + tokenizer = Tokenizer(tok_func=BaseTokenizer, + lang=self.arch.lang, + pre_rules=moses_preproc, + post_rules=[]) + return dict(processor=self._default_processor(tokenizer)) + + def _get_processor_moses_fastai(self, ds_uses_moses, add_open_file_processor=False): + moses_preproc = [MosesPreprocessingFunc(self.arch.lang)] if not ds_uses_moses else [] + tokenizer = Tokenizer(tok_func=BaseTokenizer, + lang=self.arch.lang, + pre_rules=moses_preproc + defaults.text_pre_rules, + post_rules=defaults.text_post_rules) + return dict(processor=self._default_processor(tokenizer)) + + def _get_processor_pure_fastai(self, ds_uses_moses, add_open_file_processor=False): + if not ds_uses_moses: + warn("Make sure your base model was not pretrained on moses tokenized Wikipedia (default for multifit).") + tokenizer = Tokenizer(tok_func=SpacyTokenizer, lang=self.arch.lang) + return dict(processor=self._default_processor(tokenizer)) + + def cleanup(self): + if self.temp_dir is not None: + self.temp_dir.cleanup() + self.temp_dir = None + + def __del__(self): + self.cleanup() - def _get_processor_pure_fastai(self, ds_need_moses, add_open_file_processor=False): - if ds_need_moses: - warn("fastai dont use moses, make sure you pretrained from wikpiedia that wasn't tokenized with moses.") - return dict() diff --git a/ulmfit/training.py b/ulmfit/training.py index 8c67429..c99f33b 100644 --- a/ulmfit/training.py +++ b/ulmfit/training.py @@ -1,16 +1,26 @@ +from pathlib import Path + +import torch import dataclasses from fastai.callbacks import CSVLogger, SaveModelCallback from fastai.text import * -import torch -from ulmfit.datasets import ULMFiTDataset -from pathlib import Path +from ulmfit.datasets import ULMFiTDataset,ULMFiTTokenizer CLS_BEST = 'cls_best' LM_BEST = "lm_best" ENC_BEST = "enc_best" +def detect_lang_from_dataset_path(dataset_path:Path): + lang, size = dataset_path.name.split('-') + if lang == "wikitext": + lang = "en" + if len(lang) == 2: + return lang + return None + + @dataclass class Params: def replace_(self, **changes): @@ -27,6 +37,7 @@ class Params: class ULMFiTArchitecture(Params): tokenizer: str = "f" max_vocab: int = 60000 + lang: str = None emb_sz: int = awd_lstm_lm_config['emb_sz'] n_hid: int = awd_lstm_lm_config['n_hid'] @@ -44,12 +55,12 @@ class ULMFiTArchitecture(Params): tokenizer_prefix = f"{self.tokenizer}{self.max_vocab // 1000}k" return f'models/{tokenizer_prefix}' - def dataset(self, dataset_path_or_object, **args): + def dataset(self, dataset_path_or_object, tokenizer, **args): if hasattr(dataset_path_or_object, 'load_lm_databunch'): return dataset_path_or_object if dataset_path_or_object is None: return None - return ULMFiTDataset(dataset_path=Path(dataset_path_or_object), tokenizer=self.tokenizer, max_vocab=self.max_vocab, **args) + return ULMFiTDataset(dataset_path=Path(dataset_path_or_object), tokenizer=tokenizer, **args) def set_seed(seed, name): @@ -92,15 +103,22 @@ class ULMFiTTrainingCommand(Params): def info_json(self): return self.__class__.__name__.lower().replace("ulmfit", "") + ".json" - - def _set_dataset_(self, dataset_or_path): - dataset_or_path = self.arch.dataset(dataset_or_path or self.dataset_path or getattr(self, 'base', self).dataset_path) - self.dataset_path = dataset_or_path.dataset_path - return dataset_or_path + def _set_dataset_(self, dataset_or_path, tokenizer): + #TODO: refactor, this bit is unclear (set_dataset that does nothing when is None passed?) + dataset_or_path = dataset_or_path or self.dataset_path or getattr(self, 'base', self).dataset_path + dataset = self.arch.dataset(dataset_or_path, tokenizer=tokenizer) + self.dataset_path = dataset.dataset_path + return dataset @property def dataset(self): - return self.arch.dataset(self.dataset_path) + return self.arch.dataset(self.dataset_path, self.tokneizer) + + @property + def tokenizer(self): + if self.experiment_path is None: + raise ValueError("There is no pretrained tokenizer, experiment_path is None") + return ULMFiTTokenizer(arch=self.arch, pretrained_path=self.experiment_path) def save_paramters(self): params = dataclasses.asdict(self) @@ -131,6 +149,9 @@ class ULMFiTTrainingCommand(Params): if update_arch: self.arch.replace_(**arch) self.replace_(**d) + # compatiblity with older info.json formats where lang was not stored + if self.arch.lang is None and 'dataset_path' in d: + self.arch.lang = detect_lang_from_dataset_path(Path(d['dataset_path'])) self.name = experiment_path.name dataset_path = experiment_path.parent.parent.parent # ./de-1/models/fsp15k/multfit_fp16 -> ./de-1 self.dataset_path = Path(dataset_path) @@ -183,24 +204,35 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): learn.unfreeze() learn.fit_one_cycle(self.num_epochs, self.lr, (0.8, 0.7)) - def train_(self, dataset_or_path=None, **train_config): - dataset = self._set_dataset_(dataset_or_path) + def train_(self, dataset_or_path, **train_config): + if self.arch.lang is None: + lang = detect_lang_from_dataset_path(Path(dataset_or_path)) + if lang is None: + warn("Unable to detect language from dataset path assuming English, use replace_(lang='??') change it.") + lang = 'en' + self.arch.lang = lang self.replace_(**train_config, _strict=True) set_seed(self.seed, "LM weights seed") if hasattr(self, 'base'): - dataset.use_base_model_subword_vocabulary(self.base.experiment_path) + base_tokenizer = self.base.tokenizer + else: + base_tokenizer = ULMFiTTokenizer(arch=self.arch, pretrained_path=None) + + dataset = self._set_dataset_(dataset_or_path, base_tokenizer) learn = self._learner(data_lm=dataset.load_lm_databunch(bs=self.bs, bptt=self.bptt)) experiment_path = learn.path / learn.model_dir print("Experiment", experiment_path) if self.num_epochs > 0: self._fit_schedule(learn) + self.experiment_path = experiment_path + base_tokenizer.save(self.experiment_path, learn=learn) learn.to_fp32() learn.save_encoder(ENC_BEST) learn.save(LM_BEST, with_opt=False) learn.destroy() - print("Language model saved to", self.experiment_path) self.save_paramters() + print("Language model saved to", self.experiment_path) def validate(self): raise NotImplementedError("The validation on the language model is not implemented.") @@ -305,14 +337,17 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): return learn def train_(self, dataset_or_path=None, **train_config): - dataset = self._set_dataset_(dataset_or_path) + self.replace_(**train_config, _strict=True) - dataset.use_base_model_subword_vocabulary(self.base.experiment_path) + + base_tokenizer = self.base.tokenizer + dataset = self._set_dataset_(dataset_or_path, base_tokenizer) learn = self._learner(data_clas=dataset.load_clas_databunch(bs=self.bs)) self._fit_schedule(learn) self.experiment_path = learn.path / learn.model_dir + base_tokenizer.save(self.experiment_path, learn=learn) learn.to_fp32() learn.save(CLS_BEST, with_opt=False) print("Classifier model saved to", self.experiment_path)