diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 4ac57aa..3325ab5 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -4,6 +4,8 @@ expected to have been tokenized with Moses and processed with `postprocess_wikit That is, the data is expected to be white-space separated and numbers are expected to be split. """ +from dataclasses import InitVar + import fastai import fire @@ -11,8 +13,8 @@ from fastai import * from fastai.text import * import torch from fastai_contrib.utils import read_file, read_whitespace_file, \ - validate, PAD, UNK, get_sentencepiece -from fastai_contrib.learner import bilm_learner, accuracy_fwd, accuracy_bwd + validate, PAD, UNK, get_sentencepiece, read_clas_data, TRN, VAL, TST, PAD_TOKEN_ID +from fastai_contrib.learner import bilm_learner, accuracy_fwd, accuracy_bwd, bilm_text_classifier_learner import pickle from pathlib import Path @@ -42,109 +44,153 @@ import fastai_contrib.data as contrib_data @dataclass -class Experiment: - dir_path: str +class LMHyperParams: + dataset_path: str # data_dir + + base_lm_path: str = None bidir: bool =False + qrnn: bool = True + max_vocab: int = 60000 + subword: bool = False + + emb_sz:int = 400 + nh: int = None + nl: int = 3 + + # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103) + # for training on smaller datasets, more dropout is necessary + drop_mult = 0.1 + dps = [0.25, 0.1, 0.2, 0.02, 0.15] + clip: float = 0.12 bptt: int = 70 bs: int = 70 - lang: str = 'en' - max_vocab: int = 60000 - name: str = 'wt-103' - subword: bool = False - ds_pct: float = 1.0 - qrnn: bool = True - cuda_id:int = 0 - def __post_init__(self): - self.results = {} + lang: str = 'en' + name: str = '' + cuda_id: InitVar[int] = 0 + + def __post_init__(self, cuda_id): if not torch.cuda.is_available(): print('CUDA not available. Setting device=-1.') cuda_id = -1 - torch.cuda.set_device(self.cuda_id) + torch.cuda.set_device(cuda_id) + self.dataset_path = Path(self.dataset_path) + self.base_lm_path = Path(self.base_lm_path) if self.base_lm_path is not None else None - self.dir_path = Path(self.dir_path) - assert self.dir_path.exists() - self.model_dir = self.dir_path / 'models' # removed from params, as it is absolute models location in train_clas and here it is relative - self.model_dir.mkdir(exist_ok=True) + assert self.dataset_path.exists() + self.cache_dir = self.dataset_path / 'models' / self.tok_name + self.model_dir = self.cache_dir / self.full_name + + self.model_dir.mkdir(exist_ok=True, parents=True) print('Batch size:', self.bs) print('Max vocab:', self.max_vocab) + print('Cache dir:', self.cache_dir) + print('Model dir:', self.model_dir) + self.dps = np.array(self.dps) + if self.nh is None: self.nh = 1550 if self.qrnn else 1150 - if self.qrnn: - print('Using QRNNs...') + @classmethod + def based_on(cls, base_lm_path, dataset_path, **kwargs) -> 'LMHyperParams': + with open(base_lm_path/'info.json', 'r') as f: d = json.load(f) + d['dataset_path'] = dataset_path + d['base_lm_path'] = base_lm_path + d.update(kwargs) + return cls(**d) + + def save_info(self): + from dataclasses import asdict + vals = {k: (str(v) if isinstance(v, Path) else v) for k,v in asdict(self).items()} + vals.pop('name', None) + vals.pop('lang', None) + with (self.model_dir / 'info.json').open("w") as fp: json.dump(vals, fp) + print("Saving info", self.model_dir / 'info.json') @property - def model_name(self): - return 'qrnn' if self.qrnn else 'lstm' + def tok_name(self): + pref = 'sp' if self.subword else 'v' + voc_size = self.max_vocab // 1000 + return f"{pref}{voc_size}k" - def train_lm(self, num_epochs=10): - data_lm = self.load_data() - exe = Executor(self.create_lm_learner(data_lm), exp=self) - if num_epochs > 0: - exe.learn.fit_one_cycle(num_epochs, 5e-3, (0.8, 0.7), wd=1e-7) - exe.validate() - exe.save() - return exe + @property + def full_name(self): return f"{self.model_name}_{self.name}.m" + + # todo rework + @property + def model_name(self): return ('bi' if self.bidir else '') + ('qrnn' if self.qrnn else 'lstm') + + @property + def pretrained_fnames(self): return [self.base_lm_path / 'lm_best', self.base_lm_path / '../itos'] if self.base_lm_path else None + + def train_lm(self, num_lm_epochs=10, data_lm=None): + data_lm = self.load_wiki_data() if data_lm is None else data_lm + learn = self.create_lm_learner(data_lm) + + if num_lm_epochs > 0: + if self.pretrained_fnames : + learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7)) + learn.unfreeze() + if num_lm_epochs > 0: learn.fit_one_cycle(num_lm_epochs, 1e-3, moms=(0.8, 0.7)) + else: + try: + learn.load("lm_best") + print("Weights loaded") + except FileNotFoundError: + print("Starting from random weights") + learn.fit_one_cycle(num_lm_epochs, 5e-3, (0.8, 0.7), wd=1e-7) + opt_state_path = self.model_dir / 'opt_state.pth' + print(f"Saving optimiser state at {opt_state_path}") + torch.save(learn.opt.opt.state_dict(), opt_state_path) + learn.save_encoder("enc_best") + learn.save("lm_best", with_opt=False) + print(learn.path) + + self.save_info() + return learn def create_lm_learner(self, data_lm): - # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103) - # for training on smaller datasets, more dropout is necessary - if self.qrnn: - emb_sz, nh, nl = 400, 1550, 3 - # dps = np.array([0.0, 0.0, 0.0, 0.0, 0.0]) - dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) - drop_mult = 0.1 - else: - emb_sz, nh, nl = 400, 1150, 3 - # emb_sz, nh, nl = 400, 1150, 3 - dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15]) - drop_mult = 0.1 - fastai.text.learner.default_dropout['language'] = dps + fastai.text.learner.default_dropout['language'] = self.dps lm_learner = bilm_learner if self.bidir else language_model_learner - learn = lm_learner(data_lm, bptt=self.bptt, emb_sz=emb_sz, nh=nh, nl=nl, pad_token=1, - drop_mult=drop_mult, tie_weights=True, model_dir=self.model_dir.name, - bias=True, qrnn=self.qrnn, clip=0.12) + + learn = lm_learner(data_lm, bptt=self.bptt, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, pad_token=PAD_TOKEN_ID, + drop_mult=self.drop_mult, tie_weights=True, model_dir= self.model_dir.relative_to(data_lm.path), + bias=True, qrnn=self.qrnn, clip=self.clip, pretrained_fnames=self.pretrained_fnames) # compared to standard Adam, we set beta_1 to 0.8 learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99)) learn.true_wd = False print("true_wd: ", learn.true_wd) - if self.bidir: - learn.metrics = [accuracy_fwd, accuracy_bwd] - else: - learn.metrics = [accuracy] + learn.metrics = [accuracy_fwd, accuracy_bwd] if self.bidir else [accuracy] return learn + @property + def lm_type(self): + return contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM - def load_data(self): - trn_path = self.dir_path / f'{self.lang}.wiki.train.tokens' - val_path = self.dir_path / f'{self.lang}.wiki.valid.tokens' - tst_path = self.dir_path / f'{self.lang}.wiki.test.tokens' + def load_wiki_data(self): + trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens' + val_path = self.dataset_path / f'{self.lang}.wiki.valid.tokens' + tst_path = self.dataset_path / f'{self.lang}.wiki.test.tokens' for path_ in [trn_path, val_path, tst_path]: assert path_.exists(), f'Error: {path_} does not exist.' if self.subword: # apply sentencepiece tokenization - trn_path = self.dir_path / f'{self.lang}.wiki.train.tokens' - val_path = self.dir_path / f'{self.lang}.wiki.valid.tokens' + trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens' + val_path = self.dataset_path / f'{self.lang}.wiki.valid.tokens' read_file(trn_path, 'train') read_file(val_path, 'valid') - sp = get_sentencepiece(self.dir_path, trn_path, self.name, vocab_size=self.max_vocab) + sp = get_sentencepiece(self.dataset_path, trn_path, self.name, vocab_size=self.max_vocab) lm_type = contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM - data_lm = TextLMDataBunch.from_csv(self.dir_path, 'train.csv', **sp, bs=self.bs, bptt=self.bptt, lm_type=lm_type) + data_lm = TextLMDataBunch.from_csv(self.dataset_path, 'train.csv', **sp, bs=self.bs, bptt=self.bptt, lm_type=lm_type) 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) - if self.ds_pct < 1.0: - trn_tok = trn_tok[:max(20, int(len(trn_tok) * self.ds_pct))] - val_tok = val_tok[:max(20, int(len(val_tok) * self.ds_pct))] - print(f"Limiting data sets to {self.ds_pct * 100}%, trn {len(trn_tok)}, val: {len(val_tok)}") - - itos_fname = self.model_dir / f'itos_{self.name}.pkl' + itos_fname = self.cache_dir / f'itos.pkl' if not itos_fname.exists(): # create the vocabulary cnt = Counter(word for sent in trn_tok for word in sent) @@ -154,7 +200,6 @@ class Experiment: # save vocabulary print(f"Saving vocabulary as {itos_fname}") - self.results['itos_fname'] = itos_fname with open(itos_fname, 'wb') as f: pickle.dump(itos, f) else: @@ -166,50 +211,27 @@ class Experiment: 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]) - lm_type = contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM + # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos)) - data_lm = TextLMDataBunch.from_ids(path=self.dir_path, vocab=vocab, train_ids=trn_ids, + data_lm = TextLMDataBunch.from_ids(path=self.dataset_path, vocab=vocab, train_ids=trn_ids, valid_ids=val_ids, bs=self.bs, bptt=self.bptt, - lm_type=lm_type) + lm_type=self.lm_type) 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)])) return data_lm - -class Executor: - def __init__(self, learn, exp): - self.exp = exp - self.learn = learn - try: - self.learn.load(f'{self.exp.model_name}_{self.exp.name}') - print("Weights loaded") - except FileNotFoundError: - print("Starting from random weights") - pass - - def validate(self): - if not self.exp.subword and self.exp.max_vocab is None: - raise NotImplementedError("figure out how to validate and save results") - # 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]) - logloss, perplexity = validate(learn.model, tst_ids, self.exp.bptt) - print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item()) - - def save(self): - print(f"Saving models at {self.learn.path / self.learn.model_dir}") - self.learn.save(f'{self.exp.model_name}_{self.exp.name}') - - opt_state_path = self.learn.path / self.learn.model_dir / f'{self.exp.model_name}_{self.exp.name}_state.pth' - print(f"Saving optimiser state at {opt_state_path}") - torch.save(self.learn.opt.opt.state_dict(), opt_state_path) - - self.exp.results['accuracy'] = self.learn.validate()[1] #TODO rewrite - +def validate_lm(self): + if not self.exp.subword and self.exp.max_vocab is None: + raise NotImplementedError("figure out how to validate and save results") + # 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]) + logloss, perplexity = validate(learn.model, tst_ids, self.exp.bptt) + print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item()) if __name__ == '__main__': - fire.Fire(Experiment) + fire.Fire(LMHyperParams)