diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index f92fff6..6b18d18 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -12,11 +12,21 @@ import csv import pathlib import tarfile from sklearn import model_selection +from sacremoses import MosesTokenizer +from typing import Dict, Tuple, List EOS = '' UNK = '' PAD = '' +SEP = '' # special separator token for NLI PAD_TOKEN_ID = 1 +IMDB, XNLI, TRN, VAL, TST, EN = 'imdb', 'xnli', 'train', 'val', 'test', 'en' +DATASETS = ['imdb', 'xnli'] +XNLI_PATHS = { + TRN: 'XNLI-MT-1.0/multinli/multinli.train.%s.tsv', + VAL: 'XNLI-1.0/xnli.dev.tsv', + TST: 'XNLI-1.0/xnli.test.tsv' +} number_match_re = re.compile(r'^([0-9]+[,.]?)+$') number_split_re = re.compile(r'([,.])') @@ -97,6 +107,104 @@ def prepare_imdb(file_path: str, prepare_lm = False): df_val.to_csv(LM_PATH / 'test.csv', header=False, index=False) +def read_imdb(dir_path, lang, split) -> Tuple[List[List[str]], List[str]]: + """ + Reads IMDb data. + :param dir_path: the path to the imdb folder + :param lang: the language (not used here as IMDb is only available in English) + :param split: the split of the data that should be read (train, test, val) + :return: a tuple consisting of a list of lists of tokens and a list of labels + """ + file_path = dir_path / 'train.csv' if split == TRN else dir_path / 'test.csv' + toks, lbls = [], [] + mt = MosesTokenizer('en') + print(f'Reading {file_path}...') + with open(file_path, encoding='utf-8') as f: + reader = csv.reader(f) + for row in reader: + label, text = row + lbls.append(label) + raw_tokens = mt.tokenize(text, return_str=True).split(' ') + [EOS] + tokens = [] + for token in raw_tokens: + if number_match_re.match(token): + tokens += number_split_re.sub(r' @\1@ ', token).split() + else: + tokens.append(token) + toks.append(tokens) + return toks, lbls + + +def read_xnli(dir_path, lang, split) -> Tuple[List[List[str]], List[str]]: + """ + Reads XNLI data. + :param dir_path: the path to the xnli folder + :param lang: the language + :param split: the split of the data that should be read (train, test, val) + :return: a tuple consisting of a list of lists of tokens and a list of labels + """ + file_path = XNLI_PATHS[split] + if split == TRN: + file_path = file_path % lang + elif lang == EN: + file_name = 'xnli.dev.en.tsv' if split == VAL else 'xnli.test.en.tsv' + file_path = f'XNLI-MT-1.0/xnli/{file_name}' + file_path = dir_path / file_path + toks, lbls = [], [] + print(f'Reading {file_path}...') + with open(file_path, encoding='utf-8') as f: + reader = csv.reader(f, delimiter='\t', quoting=csv.QUOTE_NONE) + for i, row in enumerate(reader): + if i == 0: # skip the header + continue + # the examples are already tokenized with Moses + if split == TRN: + premise, hypo, label = row + else: + ex_lang = row[0] + if ex_lang != lang: + continue + premise, hypo, label = row[-3], row[-2], row[1] + # TODO add BOS + premise_toks = premise.split(' ') + [EOS] + hypo_toks = hypo.split(' ') + [EOS] + toks.append(premise_toks + [SEP] + hypo_toks) + lbls.append(label) + return toks, lbls + + +def read_clas_data(dir_path, dataset, lang) -> Tuple[Dict[str, List[List[str]]], Dict[str, List[str]]]: + """ + Read the dataset from the classification datasets and tokenize them. + :param dir_path: the path to the dataset + :param dataset: the name of the dataset + :param lang: the language + :return: a tuple consisting of: + 1. a dictionary mapping splits to a list of lists of tokens + 2. a dictionary mapping splits to a list of labels + """ + processors = { + 'imdb': read_imdb, + 'xnli': read_xnli + } + processor = processors[dataset] + + toks, lbls = {}, {} + toks[TRN], lbls[TRN] = processor(dir_path, lang, TRN) + toks[TST], lbls[TST] = processor(dir_path, lang, TST) + + if dataset == IMDB: + # for IMDb, we need to split off a separate validation set + # note that we train and fine-tune ULMFiT on the full training set in the paper + # to do this, we can just keep the training set the same + trn_len = int(len(toks[TRN]) * 0.9) + toks[TRN], toks[VAL] = toks[TRN][:trn_len], toks[TRN][trn_len:] + lbls[TRN], lbls[VAL] = lbls[TRN][:trn_len], lbls[TRN][trn_len:] + else: + toks[VAL], lbls[VAL] = processor(dir_path, lang, VAL) + return toks, lbls + + def replace_number(token): """Replaces a number and returns a list of one or multiple tokens.""" if number_match_re.match(token): @@ -124,23 +232,6 @@ def read_whitespace_file(filepath): return np.array(tokens) -def read_imdb(file_path, mt): - toks, lbls = [], [] - print(f'Reading {file_path}...') - with open(file_path, encoding='utf-8') as f: - reader = csv.reader(f) - for row in reader: - label, text = row - lbls.append(label) - raw_tokens = mt.tokenize(text, return_str=True).split(' ') + [EOS] - tokens = [] - for token in raw_tokens: - if number_match_re.match(token): - tokens += number_split_re.sub(r' @\1@ ', token).split() - else: - tokens.append(token) - toks.append(tokens) - return np.array(toks), np.array(lbls) class DataStump: diff --git a/prepare_xnli.sh b/prepare_xnli.sh new file mode 100644 index 0000000..fa67589 --- /dev/null +++ b/prepare_xnli.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Script to download a Wikipedia dump + +# Script is partially based on https://github.com/facebookresearch/fastText/blob/master/get-wikimedia.sh +ROOT="data" +XNLI_DIR="${ROOT}/xnli" +mkdir -p "${ROOT}" +mkdir -p "${XNLI_DIR}" + +echo "Saving data in ""$ROOT" +MT_FILE="XNLI-MT-1.0.zip" +XNLI_FILE="XNLI-1.0.zip" +MT_PATH="${XNLI_DIR}/${MT_FILE}" +XNLI_PATH="${XNLI_DIR}/${XNLI_FILE}" + +if [ ! -f "${MT_PATH}" ]; then + wget -c "https://s3.amazonaws.com/xnli/XNLI-MT-1.0.zip" -P "${XNLI_DIR}" + wget -c "https://s3.amazonaws.com/xnli/XNLI-1.0.zip" -P "${XNLI_DIR}" +else + echo "${MT_PATH} already exists. Skipping download." +fi + +unzip "${MT_PATH}" -d "${XNLI_DIR}" +unzip "${XNLI_PATH}" -d "${XNLI_DIR}" diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py index a7a7878..f0273bf 100644 --- a/ulmfit/train_clas.py +++ b/ulmfit/train_clas.py @@ -5,85 +5,103 @@ 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, prepare_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 -CLASSES = ['neg', 'pos', 'unsup'] +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) -def new_train_clas(dir_path, lang='en', pretrain_name='wt103', model_dir='models', qrnn=False, - fine_tune=False, clean=False, max_vocab=30000, bs=70, bptt=70): - dir_path = Path(dir_path) - model_dir = dir_path / model_dir - assert dir_path.exists(), f'Error: {dir_path} does not exist.' + 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 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 as 10% of training 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_csv(tmp_dir, '...') - data_clas = TextClasDataBunch.from_tokens(tmp_dir, '...') + 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, csv_name='train.csv', test='test.csv', bs=bs, - classes=CLASSES) - data_clas = TextClasDataBunch.from_csv(dir_path, csv_name='train.csv', - vocab=data_lm.train_ds.vocab, bs=bs,classes=CLASSES) + 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 implement 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 @@ -116,6 +134,9 @@ def new_train_clas(dir_path, lang='en', pretrain_name='wt103', model_dir='models 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) diff --git a/ulmfit/xnli/prepare_xnli.sh b/ulmfit/xnli/prepare_xnli.sh deleted file mode 100644 index 31e1f0b..0000000 --- a/ulmfit/xnli/prepare_xnli.sh +++ /dev/null @@ -1,48 +0,0 @@ -#!/usr/bin/env bash -# Script to download a Wikipedia dump - -# Script is partially based on https://github.com/facebookresearch/fastText/blob/master/get-wikimedia.sh -ROOT="data" -DUMP_DIR="${ROOT}/xnli_dnld" -EXTR_DIR="${ROOT}/xnli_extr" -WIKI_DIR="${ROOT}/xnli" -EXTR="wikiextractor" -mkdir -p "${ROOT}" -mkdir -p "${DUMP_DIR}" -mkdir -p "${WIKI_DIR}" - -echo "Saving data in ""$ROOT" -read -r -p "Choose a language (e.g. en, fr, etc.): " choice -LANG="$choice" -echo "Chosen language: ""$LANG" -MT_FILE="XNLI-MT-1.0.zip" -XNLI_FILE="XNLI-1.0.zip" -MT_PATH="${DUMP_DIR}/${MT_FILE}" -XNLI_PATH="${DUMP_DIR}/${XNLI_FILE}" - -if [ ! -f "${MT_PATH}" ]; then - read -r -p "Continue to download (WARNING: This might be big and can take a long time!) (y/n)? " choice - case "$choice" in - y|Y ) echo "Starting download...";; - n|N ) echo "Exiting";exit 1;; - * ) echo "Invalid answer";exit 1;; - esac - wget -c "https://s3.amazonaws.com/xnli/XNLI-MT-1.0.zip" -P "${DUMP_DIR}" - wget -c "https://s3.amazonaws.com/xnli/XNLI-1.0.zip" -P "${DUMP_DIR}" - -else - echo "${MT_PATH} already exists. Skipping download." -fi - -if [ ! -d "${EXTR_DIR}" ]; then - read -r -p "Continue to extract XNLI (WARNING: This might take a long time!) (y/n)? " choice - case "$choice" in - y|Y ) echo "Extracting ${MT_PATH} to ${EXTR_DIR}...";; - n|N ) echo "Exiting";exit 1;; - * ) echo "Invalid answer";exit 1;; - esac - unzip "${MT_PATH}" -d "${EXTR_DIR}" - unzip "${XNLI_PATH}" -d "${EXTR_DIR}" -else - echo "${EXTR_DIR} already exists. Skipping extraction." -fi