From 94fa1a3ecbeed383a37bb53f7e49121cfbf4125b Mon Sep 17 00:00:00 2001 From: Sebastian Date: Thu, 15 Nov 2018 14:36:50 +0000 Subject: [PATCH] Encapsulated data reading in utils method, removed fastai processing, added doc string --- ulmfit/train_clas.py | 139 +++++++++++++++++++++++++------------------ 1 file changed, 82 insertions(+), 57 deletions(-) diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py index caa71be..4e4d96c 100644 --- a/ulmfit/train_clas.py +++ b/ulmfit/train_clas.py @@ -5,81 +5,102 @@ Optionally fine-tune LM before. import numpy as np import pickle +import torch from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner from fastai import fit_one_cycle -from fastai_contrib.utils import PAD, UNK, read_imdb, PAD_TOKEN_ID +from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST +from fastai.text.transform import Vocab -from sacremoses import MosesTokenizer import fire from collections import Counter from pathlib import Path -def new_train_clas(dir_path, lang='en', pretrain_name='wt-103', model_dir='models', qrnn=True, - fine_tune=True, clean=True, max_vocab=30000, bs=70, bptt=70): - dir_path = Path(dir_path) +def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt-103', model_dir='models', qrnn=True, + fine_tune=True, max_vocab=30000, bs=70, bptt=70, name='imdb-clas', + dataset='imdb'): + """ + :param data_dir: The path to the `data` directory + :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 pretrain_name: name of the pretrained model + :param model_dir: The path to the directory where the pretrained model is saved + :param qrrn: Use a QRNN. Requires installing cupy. + :param fine_tune: Fine-tune the pretrained language model + :param max_vocab: The maximum size of the vocabulary. + :param bs: The batch size. + :param bptt: The back-propagation-through-time sequence length. + :param name: The name used for both the model and the vocabulary. + :param dataset: The dataset used for evaluation. Currently only IMDb and + XNLI are implemented. Assumes dataset is located in `data` + folder and that name of folder is the same as dataset name. + """ + if not torch.cuda.is_available(): + print('CUDA not available. Setting device=-1.') + cuda_id = -1 + torch.cuda.set_device(cuda_id) + + print(f'Dataset: {dataset}. Language: {lang}.') + assert dataset in DATASETS, f'Error: {dataset} processing is not implemented.' + assert (dataset == 'imdb' and lang == 'en') or not dataset == 'imdb',\ + 'Error: IMDb is only available in English.' + + data_dir = Path(data_dir) + assert data_dir.name == 'data',\ + f'Error: Name of data directory should be data, not {data_dir.name}.' + dataset_dir = data_dir / dataset model_dir = Path(model_dir) - assert dir_path.exists(), f'Error: {dir_path} does not exist.' + assert data_dir.exists(), f'Error: {data_dir} does not exist.' + assert dataset_dir.exists(), f'Error: {dataset_dir} does not exist.' assert model_dir.exists(), f'Error: {model_dir} does not exist.' if qrnn: print('Using QRNNs...') + model_name = 'qrnn' if qrnn else 'lstm' - if clean: - # use no preprocessing besides MosesTokenizer - tmp_dir = dir_path / 'tmp' - tmp_dir.mkdir(exist_ok=True) - if not (tmp_dir / 'train_ids.npy').exists(): - trn_path = dir_path / 'train.csv' - tst_path = dir_path / 'test.csv' - assert trn_path.exists(), f'Error: {trn_path} does not exist.' - assert tst_path.exists(), f'Error: {tst_path} does not exist.' - trn_toks, trn_lbls = read_imdb(trn_path, MosesTokenizer(lang)) - tst_toks, tst_lbls = read_imdb(tst_path, MosesTokenizer(lang)) + tmp_dir = dataset_dir / 'tmp' + tmp_dir.mkdir(exist_ok=True) + vocab_file = tmp_dir / f'vocab_{lang}.pkl' - # split off validation set if it does not exist - val_path = dir_path / 'valid.csv' - if not val_path.exists(): - trn_len = int(len(trn_toks) * 0.9) - trn_toks, val_toks = trn_toks[:trn_len], trn_toks[trn_len:] - trn_lbls, val_lbls = trn_lbls[:trn_len], trn_lbls[trn_len:] - else: - val_toks, val_lbls = read_imdb(val_path, MosesTokenizer(lang)) + if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists(): + print('Reading the data...') + toks, lbls = read_clas_data(dataset_dir, dataset, lang) - # create the vocabulary - cnt = Counter(word for example in trn_toks for word in example) - itos = [o for o, c in cnt.most_common(n=max_vocab)] - itos.insert(0, PAD) - itos.insert(0, UNK) - stoi = {w: i for i, w in enumerate(itos)} - with open(tmp_dir / 'itos.pkl', 'wb') as f: - pickle.dump(itos, f) + # create the vocabulary + counter = Counter(word for example in toks[TRN] for word in example) + itos = [word for word, count in counter.most_common(n=max_vocab)] + itos.insert(0, PAD) + itos.insert(0, UNK) + vocab = Vocab(itos) + stoi = vocab.stoi + with open(vocab_file, 'wb') as f: + pickle.dump(vocab, f) - trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_toks]) - val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_toks]) - tst_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in tst_toks]) - - print(f'Train size: {len(trn_ids)}. Valid size: {len(val_ids)}. ' - f'Test size: {len(tst_ids)}.') - - for split, ids, lbl in zip(['train', 'valid', 'test'], - [trn_ids, val_ids, tst_ids], - [trn_lbls, val_lbls, tst_lbls]): - np.save(tmp_dir / f'{split}_ids.npy', ids) - np.save(tmp_dir / f'{split}_lbl.npy', lbl) - - data_lm = TextLMDataBunch.from_id_files(tmp_dir, test='test') - data_clas = TextClasDataBunch.from_id_files(tmp_dir, test='test') + ids = {} + for split in [TRN, VAL, TST]: + ids[split] = np.array([([stoi.get(w, stoi[UNK]) for w in s]) + for s in toks[split]]) + np.save(tmp_dir / f'{split}_{lang}_ids.npy', ids[split]) + np.save(tmp_dir / f'{split}_{lang}_lbl.npy', lbls[split]) else: - # use fastai peprocessing and tokenization - data_lm = TextLMDataBunch.from_csv(dir_path, bs=bs) - data_clas = TextClasDataBunch.from_csv(dir_path, vocab=data_lm.train_ds.vocab, bs=bs) + print('Loading the pickled data...') + ids, lbls = {}, {} + for split in [TRN, VAL, TST]: + ids[split] = np.load(tmp_dir / f'{split}_{lang}_ids.npy') + lbls[split] = np.load(tmp_dir / f'{split}_{lang}_lbl.npy') + with open(vocab_file, 'rb') as f: + vocab = pickle.load(f) - # todo implemend save and load - #data_lm.save() - #data_clas.save() - #data_lm = TextLMDataBunch.load(path) - #data_clas = TextClasDataBunch.load(path, bs=bs) + print(f'Train size: {len(ids[TRN])}. Valid size: {len(ids[VAL])}. ' + f'Test size: {len(ids[TST])}.') + + data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, trn_ids=ids[TRN], + val_ids=ids[VAL], bs=bs, bptt=bptt) + # TODO TextClasDataBunch allows tst_ids as input, but not tst_lbls? + data_clas = TextClasDataBunch.from_ids( + path=tmp_dir, vocab=vocab, trn_ids=ids[TRN], val_ids=ids[VAL], + trn_lbls=lbls[TRN], val_lbls=lbls[VAL], bs=bs) if qrnn: emb_sz, nh, nl = 400, 1550, 3 @@ -112,5 +133,9 @@ def new_train_clas(dir_path, lang='en', pretrain_name='wt-103', model_dir='model learn.unfreeze() fit_one_cycle(learn, 10, 5e-3, (0.8, 0.7), wd=1e-7) + print(f"Saving models at {learn.path / learn.model_dir}") + learn.save(f'{model_name}_{name}') -if __name__ == '__main__': fire.Fire(new_train_clas) + +if __name__ == '__main__': + fire.Fire(new_train_clas)