diff --git a/ulmfit/datasets/__init__.py b/ulmfit/datasets/__init__.py index c2a3850..3f2574e 100644 --- a/ulmfit/datasets/__init__.py +++ b/ulmfit/datasets/__init__.py @@ -11,8 +11,8 @@ from fastai.text import * import torch from ulmfit.datasets.utils import read_whitespace_file, \ validate, UNK -from fastai_contrib.text_data import MosesPreprocessingFunc, get_sentencepiece_fastai, \ - make_data_bunch_from_df +from fastai_contrib.text_data import MosesPreprocessingFunc, \ + make_data_bunch_from_df, SPProcessor2 import pickle from pathlib import Path @@ -54,44 +54,56 @@ class Dataset: use_tst_for_lm: bool = False noise: float = 0.0 limit: int = None - + ds_type:str = None def __post_init__(self): self.add_trn_to_lm = True self._trn_df = None self._tst_df = None self._val_df = None - if 'wiki' in str(self.dataset_path) and len(list(self.dataset_path.glob('*.wiki.*.tokens'))) >= 2: + if self.ds_type is None: + self.ds_type = str(self.dataset_path) + + if 'wiki' in self.ds_type and len(list(self.dataset_path.glob('*.wiki.*.tokens'))) >= 2: + self.ds_type = 'wiki' self._post_init_tokenized_wiki() - elif 'reddit' in str(self.dataset_path): + elif 'wiki' in self.ds_type and len(list(self.dataset_path.glob('wiki.*.tokens'))) >= 2: + self.ds_type = 'wiki' + self._post_init_tokenized_wiki(wiki103=True) + elif 'reddit' in self.ds_type: + self.ds_type = 'reddit' self._post_init_default_csv( lang='en', uses_moses=False, add_trn_to_lm=True, use_lang_as_prefix=False) - elif 'xnli' in str(self.dataset_path): + elif 'xnli' in self.ds_type: + self.ds_type = 'xnli' raise NotImplementedError("Support for XNLI is not implemented yet") - elif 'imdb' in self.dataset_path.name: + elif 'imdb' in self.ds_type: + self.ds_type = 'imdb' self._post_init_default_csv( lang='en', uses_moses=True, add_trn_to_lm=True, use_lang_as_prefix=False) - elif 'mldoc' in str(self.dataset_path): + elif 'mldoc' in self.ds_type: + self.ds_type = 'mldoc' self._post_init_default_csv( lang=self._language_from_dataset_path(), uses_moses=False, add_trn_to_lm=False, use_lang_as_prefix=True) - elif 'hate' in str(self.dataset_path): + elif 'hate' in self.ds_type: + self.ds_type = 'hate' self._post_init_default_csv( lang=self._language_from_dataset_path(), uses_moses=False, add_trn_to_lm=True, use_lang_as_prefix=True) else: - raise NotImplementedError(f"Not supported dataset {self.dataset_path}") + raise NotImplementedError(f"Not supported dataset {self.dataset_path} {self.ds_type}") def _post_init_default_csv(self, lang, uses_moses, add_trn_to_lm, use_lang_as_prefix): self.lang = lang @@ -112,18 +124,22 @@ class Dataset: self.tst_path = self.dataset_path / f'{prefix}test.csv' self.unsup_path = self.dataset_path / f'{prefix}unsup.csv' - def _post_init_tokenized_wiki(self): + def _post_init_tokenized_wiki(self, wiki103=False): self.uses_moses = True self.use_tst_for_lm = False self.add_trn_to_lm = True self.lang = self._language_from_dataset_path() self._read_data = read_wiki_articles + if wiki103: + prefix="" + else: + prefix=f"{self.lang}." - self.trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens' - self.val_path = self.dataset_path / f'{self.lang}.wiki.valid.tokens' - self.tst_path = self.dataset_path / f'{self.lang}.wiki.test.tokens' - self.unsup_path = self.dataset_path / f'{self.lang}.wiki.unsup.tokens' + self.trn_path = self.dataset_path / f'{prefix}wiki.train.tokens' + self.val_path = self.dataset_path / f'{prefix}wiki.valid.tokens' + self.tst_path = self.dataset_path / f'{prefix}wiki.test.tokens' + self.unsup_path = self.dataset_path / f'{prefix}wiki.unsup.tokens' def _language_from_dataset_path(self): lang, size = self.dataset_path.name.split('-') @@ -191,11 +207,12 @@ class Dataset: class ULMFiTDataset(Dataset): tokenizer: str = 'f' max_vocab: int = 60000 - + cache_path: Path = None def __post_init__(self): super().__post_init__() - tokenizer_prefix = f"{self.tokenizer}{self.max_vocab // 1000}k" - self.cache_path = self.dataset_path / "models" / tokenizer_prefix + 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._vocab = None def use_base_model_subword_vocabulary(self, base_lm_path: Path): @@ -204,13 +221,22 @@ class ULMFiTDataset(Dataset): 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 and (base_lm_path / '..' / 'spm.vocab').exists() and \ - (base_lm_path.parent.resolve() != self.cache_path.resolve()): - shutil.copy(str(base_lm_path / '..' / 'itos.pkl'), str(self.cache_path)) - shutil.copy(str(base_lm_path / '..' / 'spm.model'), str(self.cache_path)) - shutil.copy(str(base_lm_path / '..' / 'spm.vocab'), str(self.cache_path)) + 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 diff --git a/ulmfit/training.py b/ulmfit/training.py index 086825a..ca56ffe 100644 --- a/ulmfit/training.py +++ b/ulmfit/training.py @@ -33,7 +33,7 @@ class ULMFITArchitecture(Params): n_layers: int = awd_lstm_lm_config['n_layers'] qrnn: bool = awd_lstm_lm_config['qrnn'] - def model_name(self, name): + def model_name(self, name=""): model_suffix = '' # if self.lmseed is None else f'_lmseed-{self.lmseed}' model_prefix = 'qrnn' if self.qrnn else 'lstm' @@ -47,6 +47,8 @@ class ULMFITArchitecture(Params): def dataset(self, dataset_path_or_object, **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) @@ -84,7 +86,7 @@ class ULMFiTTrainingCommand(Params): @property def model_name(self): return (self.name or self.arch.model_name()) + ( - "" if self.seed == 0 or "seed" in self.name else f"seed{self.seed}") + "" if self.seed is None or self.seed == 0 or "seed" in self.name else f"seed{self.seed}") @property def info_json(self): @@ -149,7 +151,7 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): fp16: bool = False lr: float = 5e-3 - def _learner(self, dataset, **additional_trn_args): + def _learner(self, data_lm, **additional_trn_args): config = awd_lstm_lm_config.copy() config.update(emb_sz=self.arch.emb_sz, n_hid=self.arch.n_hid, n_layers=self.arch.n_layers, qrnn=self.arch.qrnn) @@ -157,7 +159,6 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): pretrained=False) trn_args.update(**additional_trn_args) print("Training args: ", trn_args, "config: ", config) - data_lm = dataset.load_lm_databunch(bs=self.bs, bptt=self.bptt) learn = language_model_learner(data_lm, AWD_LSTM, config=config, @@ -188,7 +189,7 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): set_seed(self.seed, "LM weights seed") if hasattr(self, 'base'): dataset.use_base_model_subword_vocabulary(self.base.experiment_path) - learn = self._learner(dataset) + 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: @@ -227,12 +228,12 @@ class ULMFiTFinetuning(ULMFiTPretraining): def __post_init__(self): self.lr = 1e-3 - def _learner(self, dataset, **additional_trn_args): + def _learner(self, data_lm, **additional_trn_args): pretrained_fnames = None if self.base is None else self.base.model_fnames - if self.pretrained and pretrained_fnames is None and dataset.lang != 'en': - warn( - "You are using fastai english langauge model for {data_lm.lang}, you might be better off with just random weights.") - return super()._learner(dataset, pretrained=self.pretrained, pretrained_fnames=pretrained_fnames, + # data_lm.lang is added after dataloading + if self.pretrained and pretrained_fnames is None and data_lm.lang != 'en': + warn("You are using fastai english langauge model for {data_lm.lang}, you might be better off with just random weights.") + return super()._learner(data_lm, pretrained=self.pretrained, pretrained_fnames=pretrained_fnames, **additional_trn_args) def _fit_schedule(self, learn): @@ -245,7 +246,6 @@ class ULMFiTFinetuning(ULMFiTPretraining): else: super()._fit_schedule(learn) - @dataclass class ULMFiTClassifier(ULMFiTTrainingCommand): bs: int = 20 @@ -263,7 +263,7 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): fp16: bool = False arch: ULMFITArchitecture = None - def _learner(self, dataset, eval_only=False, **additional_trn_args): + def _learner(self, data_clas, eval_only=False, **additional_trn_args): assert self.weighted_cross_entropy is None or self.label_smoothing_eps == 0, "Label smoohting not implemented with weighted_cross_entropy" if self.weighted_cross_entropy is not None: loss_func = CrossEntropyFlat(weight=torch.tensor(self.weighted_cross_entropy, dtype=torch.float32).cuda()) @@ -273,7 +273,6 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): loss_func = None set_seed(self.seed, "Classifier weights seed") - data_clas, data_tst = dataset.load_clas_databunch(bs=self.bs) config = awd_lstm_clas_config.copy() config.update(emb_sz=self.arch.emb_sz, n_hid=self.arch.n_hid, n_layers=self.arch.n_layers, qrnn=self.arch.qrnn) @@ -288,8 +287,8 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): model_dir=self.model_name, silent=eval_only, **trn_args) - learn.data.test_dl = data_tst.valid_dl - if self.base and not self.random_init: + + if self.base.encoder_fname and not self.random_init: print("Loading pretrained model", self.base.encoder_fname) learn.load_encoder(self.base.encoder_fname) learn.freeze() @@ -309,7 +308,7 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): dataset = self._set_dataset_(dataset_or_path) self.replace_(**train_config, _strict=True) dataset.use_base_model_subword_vocabulary(self.base.experiment_path) - learn = self._learner(dataset) + learn = self._learner(data_clas=dataset.load_clas_databunch(bs=self.bs)) self._fit_schedule(learn)