From 12e52b182e2624246651d7a95324fe43d60ddb03 Mon Sep 17 00:00:00 2001 From: Piotr Czapla Date: Thu, 8 Nov 2018 20:33:10 +0100 Subject: [PATCH] Initial import from fastaiv1 initial import from https://github.com/fastai/fastai/tree/ulmfit_v1/courses/dl2/imdb_scripts --- .gitignore | 131 +++++++++++++++++++++++ ulmfit/README.md | 186 +++++++++++++++++++++++++++++++++ ulmfit/__init__.py | 0 ulmfit/create_wikitext.py | 114 ++++++++++++++++++++ ulmfit/postprocess_wikitext.py | 121 +++++++++++++++++++++ ulmfit/prepare_wiki.sh | 52 +++++++++ ulmfit/pretrain_lm.py | 134 ++++++++++++++++++++++++ ulmfit/train_clas.py | 116 ++++++++++++++++++++ ulmfit/utils.py | 127 ++++++++++++++++++++++ 9 files changed, 981 insertions(+) create mode 100644 .gitignore create mode 100644 ulmfit/README.md create mode 100644 ulmfit/__init__.py create mode 100644 ulmfit/create_wikitext.py create mode 100644 ulmfit/postprocess_wikitext.py create mode 100644 ulmfit/prepare_wiki.sh create mode 100644 ulmfit/pretrain_lm.py create mode 100644 ulmfit/train_clas.py create mode 100644 ulmfit/utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d148f25 --- /dev/null +++ b/.gitignore @@ -0,0 +1,131 @@ +*.bak +*.log +*~ +.~* +.pypirc +~* +tmp* +tags + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +conda-dist +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# dotenv +.env + +# virtualenv +.venv +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +.vscode +*.swp + +# osx generated files +.DS_Store +.DS_Store? +.Trashes +ehthumbs.db +Thumbs.db +.idea + +# pytest +.pytest_cache + +# tools/trust-doc-nbs +examples/.last_checked + +# symlinks to fastai +examples/fastai diff --git a/ulmfit/README.md b/ulmfit/README.md new file mode 100644 index 0000000..2badf99 --- /dev/null +++ b/ulmfit/README.md @@ -0,0 +1,186 @@ +## Instructions + +### 0. Preparing Wikipedia + +If you want to train your own language model on a Wikipedia in your chosen language, +run `prepare_wiki.sh`. The script will ask for a language and will then +download, extract, and prepare the latest version of Wikipedia for the chosen language. +Note that for English (due to the size of the English Wikipedia), the extraction process +takes quite long. + +Example command: `bash prepare_wiki.sh` + +This will create a `data` folder in this directory and `wiki_dumps`, `wiki_extr`, and +`wiki` subfolders. In each subfolder, it will furthermore create a folder `LANG` +where `LANG` is the language of the Wikipedia. The prepared files are stored in +`wiki/LANG` as `train.csv` and `val.csv` to match the format used for text +classification datasets. By default, `train.csv` contains around 100 million tokens +and `val.csv` is 10% the size of `train.csv`. + +### 1. Tokenization + +Run `create_toks.py` to tokenize the input texts. + +Example command: `python create_toks.py data/imdb` + +Usage: + +``` +create_toks.py DIR_PATH [CHUNKSIZE] [N_LBLS] [LANG] +create_toks.py --dir-path DIR_PATH [--chunksize CHUNKSIZE] [--n-lbls N_LBLS] [--lang LANG] +``` + +- `DIR_PATH`: the directory where your data is located +- `CHUNKSIZE`: the size of the chunks when reading the files with pandas; use smaller sizes with less RAM +- `LANG`: the language of your corpus. + +The script expects `train.csv` and `val.csv` files to be in `DIR_PATH`. Each file should be in +CSV format. If the data is labeled, the first column should consist of the label as an integer. +The remaining columns should consist of text or features, which will be concatenated to form +each example. If the data is unlabeled, the file should just consist of a single text column. +The script will then save the training and test tokens and labels as arrays to binary files in NumPy format +in a `tmp` in the above path in the following files: +`tok_trn.npy`, `tok_val.npy`, `lbl_trn.npy`, and `lbl_val.npy`. +In addition, a joined corpus containing white space-separated tokens is produced in `tmp/joined.txt`. + +### 2. Mapping tokens to ids + +Run `tok2id.py` to map the tokens in the `tok_trn.npy` and `tok_val.npy` files to ids. + +Example command: `python tok2id.py data/imdb` + +Usage: +``` +tok2id.py PREFIX [MAX_VOCAB] [MIN_FREQ] +tok2id.py --prefix PREFIX [--max-vocab MAX_VOCAB] [--min-freq MIN_FREQ] +``` +- `PREFIX`: the file path prefix in `data/nlp_clas/{prefix}` +- `MAX_VOCAB`: the maximum vocabulary size +- `MIN_FREQ`: the minimum frequency of words that should be kept + +### (3a. Pretrain the Wikipedia language model) + +Before fine-tuning the language model, you can run `pretrain_lm.py` to create a +pre-trained language model using WikiText-103 (or whatever base corpus you prefer). + +Example command: `python pretrain_lm.py data/wiki/de/ 0 --lr 1e-3 --cl 12` + +Usage: +``` +pretrain_lm.py DIR_PATH CUDA_ID [CL] [BS] [BACKWARDS] [LR] [SAMPLED] [PRETRAIN_ID] +pretrain_lm.py --dir-path DIR_PATH --cuda-id CUDA_ID [--cl CL] [--bs BS] [--backwards BACKWARDS] [--lr LR] [--sampled SAMPLED] [--pretrain-id PRETRAIN_ID] +``` +- `DIR_PATH`: the directory that contains the Wikipedia files +- `CUDA_ID`: the id of the GPU that should be used; +- `CL`: the # of epochs to train +- `BS`: the batch size +- `BACKWARDS`: whether a backwards LM should be trained +- `LR`: the learning rate +- `SAMPLED`: whether a sampled softmax should be used (default: `True`) +- `PRETRAIN_ID`: the id used for saving the trained LM + +You might have to adapt the learning rate and the # of epochs to maximize performance. + +### 3b. Fine-tune the LM + +Alternatively, you can download the pre-trained models [here](http://files.fast.ai/models/wt103/). Before, +create a directory `wt103`. In `wt103`, create a `models` and a `tmp` folder. Save the model files +in the `models` folder and `itos_wt103.pkl`, the word-to-token mapping, to the `tmp` folder. + +Then run `finetune_lm.py` to fine-tune a language model pretrained on WikiText-103 data on the target task data. + +Example command: `python finetune_lm.py data/imdb data/wt103 1 25 --lm-id pretrain_wt103` + +Usage: +``` +finetune_lm.py DIR_PATH PRETRAIN_PATH [CUDA_ID] [CL] [PRETRAIN_ID] [LM_ID] [BS] [DROPMULT] [BACKWARDS] [LR] [PRELOAD] [BPE] [STARTAT] [USE_CLR] [USE_REGULAR_SCHEDULE] [USE_DISCRIMINATIVE] [NOTRAIN] [JOINED] [TRAIN_FILE_ID] [EARLY_STOPPING] +finetune_lm.py --dir-path DIR_PATH --pretrain-path PRETRAIN_PATH [--cuda-id CUDA_ID] [--cl CL] [--pretrain-id PRETRAIN_ID] [--lm-id LM_ID] [--bs BS] [--dropmult DROPMULT] [--backwards BACKWARDS] [--lr LR] [--preload PRELOAD] [--bpe BPE] [--startat STARTAT] [--use-clr USE_CLR] [--use-regular-schedule USE_REGULAR_SCHEDULE] [--use-discriminative USE_DISCRIMINATIVE] [--notrain NOTRAIN] [--joined JOINED] [--train-file-id TRAIN_FILE_ID] [--early-stopping EARLY_STOPPING] +``` +- `DIR_PATH`: the directory where the `tmp` and `models` folder are located +- `PRETRAIN_PATH`: the path where the pretrained model is saved; if using the downloaded model, this is `wt103` +- `CUDA_ID`: the id of the GPU used for training the model +- `CL`: number of epochs to train the model +- `PRETRAIN_ID`: the id of the pretrained model; set to `wt103` per default +- `LM_ID`: the id used for saving the fine-tuned language model +- `BS`: the batch size used for training the model +- `DROPMULT`: the factor used to multiply the dropout parameters +- `BACKWARDS`: whether a backwards LM is trained +- `LR`: the learning rate +- `PRELOAD`: whether we load a pretrained LM (`True` by default) +- `BPE`: whether we use byte-pair encoding (BPE) +- `STARTAT`: can be used to continue fine-tuning a model; if `>0`, loads an already fine-tuned LM; can also be used to indicate the layer at which to start the gradual unfreezing (`1` is last hidden layer, etc.); in the final model, we only used this for training the classifier +- `USE_CLR`: whether to use slanted triangular learning rates (STLR) (`True` by default) +- `USE_REGULAR_SCHEDULE`: whether to use a regular schedule (instead of STLR) +- `USE_DISCRIMINATIVE`: whether to use discriminative fine-tuning (`True` by default) +- `NOTRAIN`: whether to skip fine-tuning +- `JOINED`: whether to fine-tune the LM on the concatenation of training and validation data +- `TRAIN_FILE_ID`: can be used to indicate different training files (e.g. to test training sizes) +- `EARLY_STOPPING`: whether to use early stopping + +The language model is fine-tuned using warm-up reverse annealing and triangular learning rates. For IMDb, +we set `--cl`, the number of epochs to `50` and used a learning rate `--lr` of `4e-3`. + +### 4. Train the classifier + +Run `train_clas.py` to train the classifier on top of the fine-tuned language model with gradual unfreezing, +discriminative fine-tuning, and slanted triangular learning rates. + +Example command: `python train_clas.py data/imdb 0 --lm-id pretrain_wt103 --clas-id pretrain_wt103 --cl 50` + +Usage: +``` +train_clas.py DIR_PATH CUDA_ID [LM_ID] [CLAS_ID] [BS] [CL] [BACKWARDS] [STARTAT] [UNFREEZE] [LR] [DROPMULT] [BPE] [USE_CLR] [USE_REGULAR_SCHEDULE] [USE_DISCRIMINATIVE] [LAST] [CHAIN_THAW] [FROM_SCRATCH] [TRAIN_FILE_ID] +train_clas.py --dir-path DIR_PATH --cuda-id CUDA_ID [--lm-id LM_ID] [--clas-id CLAS_ID] [--bs BS] [--cl CL] [--backwards BACKWARDS] [--startat STARTAT] [--unfreeze UNFREEZE] [--lr LR] [--dropmult DROPMULT] [--bpe BPE] [--use-clr USE_CLR] [--use-regular-schedule USE_REGULAR_SCHEDULE] [--use-discriminative USE_DISCRIMINATIVE] [--last LAST] [--chain-thaw CHAIN_THAW] [--from-scratch FROM_SCRATCH] [--train-file-id TRAIN_FILE_ID] +``` +- `DIR_PATH`: the directory where the `tmp` and `models` folder are located +- `CUDA_ID`: the id of the GPU used for training the model +- `LM_ID`: the id of the fine-tuned language model that should be loaded +- `CLAS_ID`: the id used for saving the classifier +- `BS`: the batch size used for training the model +- `CL`: the number of epochs to train the model with all layers unfrozen +- `BACKWARDS`: whether a backwards LM is trained +- `STARTAT`: whether to use gradual unfreezing (`0`) or load the pretrained model (`1`) +- `UNFREEZE`: whether to unfreeze the whole network (after optional gradual unfreezing) or train only the final classifier layer (default is `True`) +- `LR`: the learning rate +- `DROPMULT`: the factor used to multiply the dropout parameters +- `BPE`: whether we use byte-pair encoding (BPE) +- `USE_CLR`: whether to use slanted triangular learning rates (STLR) (`True` by default) +- `USE_REGULAR_SCHEDULE`: whether to use a regular schedule (instead of STLR) +- `USE_DISCRIMINATIVE`: whether to use discriminative fine-tuning (`True` by default) +- `LAST`: whether to fine-tune only the last layer of the model +- `CHAIN_THAW`: whether to use chain-thaw +- `FROM_SCRATCH`: whether to train the model from scratch (without loading a pretrained model) +- `TRAIN_FILE_ID`: can be used to indicate different training files (e.g. to test training sizes) + +For fine-tuning the classifier on IMDb, we set `--cl`, the number of epochs to `50`. + +### 5. Evaluate the classifier + +Run `eval_clas.py` to get the classifier accuracy and confusion matrix. + +This requires the files produced during the training process: itos.pkl and the classifier (named clas_1.h5 by default), as well as the `npy` files containing the evaluation samples and labels. + +Example command: `python eval_clas.py data/imdb 0 --lm-id pretrain_wt103 --clas-id pretrain_wt103` + +Usage: +``` +eval_clas.py DIR_PATH CUDA_ID [LM_ID] [CLAS_ID] [BS] [BACKWARDS] [BPE] +eval_clas.py --dir-path DIR_PATH --cuda-id CUDA_ID [--lm-id LM_ID] [--clas-id CLAS_ID] [--bs BS] [--bpe BPE] +``` +- `DIR_PATH`: the directory where the `tmp` and `models` folder are located +- `CUDA_ID`: the id of the GPU used for training the model +- `LM_ID`: the id of the fine-tuned language model that should be loaded +- `CLAS_ID`: the id used for saving the classifier +- `BS`: the batch size used for training the model +- `BACKWARDS`: whether a backwards LM is trained +- `BPE`: whether we use byte-pair encoding (BPE) + +### 6. Try the classifier on text + +Run `predict_with_classifier.py` to predict against free text entry. + +This requires two files produced during the training process: the id-to-token mapping `itos.pkl` and the classifier (named `clas_1.h5` by default) + +Example command: `python predict_with_classifier.py trained_models/itos.pkl trained_models/classifier_model.h5` + +It is suggested to customize this script to your needs. diff --git a/ulmfit/__init__.py b/ulmfit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py new file mode 100644 index 0000000..5d24284 --- /dev/null +++ b/ulmfit/create_wikitext.py @@ -0,0 +1,114 @@ +""" +Script to create small and large WikiText datasets from Wikipedia articles in +any language that were downloaded with `prepare_wiki.sh`. +Articles are tokenized using the Moses tokenizer. Articles with least than +100 tokens are removed. +""" +import argparse +from pathlib import Path +import json + +from shutil import copyfile + +from sacremoses import MosesTokenizer + + +def get_texts(root): + for dir_ in root.iterdir(): + for wiki_file in dir_.iterdir(): + with open(wiki_file, encoding='utf-8') as f_in: + for line in f_in: + article = json.loads(line) + text = article['text'] + title = article['title'] + if text.strip() == title: + # print('No content continuing...') + continue + yield text + + +def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'): + total_num_tokens = 0 + print(f'Writing to {file_path}...') + i = 0 + with open(file_path, mode, encoding='utf-8') as f_out: + for i, text in enumerate(text_iter): + + num_tokens_article = 0 # count the number of tokens in an article + tokenized_paragraphs = [] + paragraphs = text.split('\n') + + for paragraph in paragraphs: + tokenized = mt.tokenize(paragraph.strip(), return_str=True) + tokenized_paragraphs.append(tokenized) + + tokens = tokenized.split(' ') # split on whitespace to keep newlines + # don't count empty lines + tokens = [token for token in tokens if token] + + # calculate length based on tokens; add 1 for newline + num_tokens_article += len(tokens) + 1 + + if num_tokens_article < 100: + # only use articles that have at least 100 tokens + continue + + for tokenized in tokenized_paragraphs: + f_out.write(tokenized + '\n') + + total_num_tokens += num_tokens_article + 1 + if total_num_tokens > num_tokens: + break + if i % 10000 == 0 and i > 0: + print('Processed {:,} documents. Total # tokens: {:,}.'.format(i, total_num_tokens)) + print('{}. # documents: {:,}. # tokens: {:,}.'.format( + file_path, i, total_num_tokens)) + + +def main(args): + + input_path = Path(args.input) + output = Path(args.output) + assert input_path.exists(), f'Error: {input_path} does not exist.' + output.mkdir(exist_ok=True) + + mt = MosesTokenizer(args.lang) + + sml_wiki = output / f'{args.lang}-2' + lrg_wiki = output / f'{args.lang}-100' + sml_wiki.mkdir(exist_ok=True) + lrg_wiki.mkdir(exist_ok=True) + + text_iter = get_texts(input_path) + + splits = ['train', 'valid', 'test'] + token_nums = [2000000, 200000, 200000] + for split, token_num in zip(splits, token_nums): + sml_file_path = sml_wiki / f'{args.lang}.wiki.{split}.tokens' + write_wikitext(sml_file_path, text_iter, mt, token_num) + lrg_file_path = lrg_wiki / f'{args.lang}.wiki.{split}.tokens' + + # copy the content of the small file to the large file + print(f'Copying {sml_file_path} to {lrg_file_path}.') + copyfile(sml_file_path, lrg_file_path) + + # add the new articles to the existing ones + lrg_wiki_train = lrg_wiki / f'{args.lang}.wiki.train.tokens' + write_wikitext(lrg_wiki_train, text_iter, mt, 98000000, mode='a') + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser() + parser.add_argument('-i', '--input', required=True, + help='the directory where the Wikipedia data extracted ' + 'with WikiExtractor.py is located. Consists of ' + 'directories AA, AB, AC, etc.') + parser.add_argument('-o', '--output', required=True, + help='the output directory where the merged Wikipedia ' + 'documents should be saved') + parser.add_argument('-l', '--lang', required=True, + help='the iso code of the language of the Wikipedia ' + 'documents, e.g. en, fr, de, etc.') + args = parser.parse_args() + main(args) diff --git a/ulmfit/postprocess_wikitext.py b/ulmfit/postprocess_wikitext.py new file mode 100644 index 0000000..3b5f595 --- /dev/null +++ b/ulmfit/postprocess_wikitext.py @@ -0,0 +1,121 @@ +""" +Script to post-process WikiText files created with `create_wikitext.py`. +Creates additional files where words not in the training data are replaced +with and numbers are modified with a regex. +""" +import argparse + +from collections import Counter +from pathlib import Path +from ulmfit.utils import replace_number, UNK + + +def build_vocab(file_path, cutoff=3): + counter = Counter() + with open(file_path, 'r', encoding='utf-8') as f: + for i, line in enumerate(f): + tokens = line.strip().split(' ') + [''] + counter.update(tokens) + vocab = {} + in_vocab_count = 0 + OOV_count = 0 + for token, count in counter.most_common(): + if count >= cutoff: + vocab[token] = count + in_vocab_count += count + else: + OOV_count += count + print('OOV ratio: %.4f.' % (OOV_count / (in_vocab_count + OOV_count))) + return vocab + + +def limit_vocab(unk_path, vocab): + """ + https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-limit_vocab-py + :param unk_path: + :param vocab: + :return: + """ + temp_file_path = unk_path.with_name(unk_path.name + '.temp') + total_num_tokens = 0 + print(f'Limiting vocab in {unk_path}. Writing to {unk_path}.') + with open(unk_path, 'r', encoding='utf-8') as f_in, open(temp_file_path, 'w', encoding='utf-8') as f_out: + for line in f_in: + tokens = [x for x in line.strip().split(' ') if x] + tokens = [token if token in vocab else UNK for token in tokens] + # Ensures there's a space between tokens, including the last word, + # newline, and the first word of the next line + tokens = tokens + ['\n'] + total_num_tokens += len(tokens) + tokens = [''] + tokens + line = ' '.join(tokens) + f_out.write(line) + print(f'{unk_path.name}. # of tokens: {total_num_tokens}') + temp_file_path.replace(unk_path) + + +def replace_numbers(file_path, unk_path): + """ + Replace numbers as in Smerity's script: + https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-post_process-py + :param file_path: + :param unk_path: + :return: + """ + print(f'Replacing numbers in {file_path}. Writing to {unk_path}.') + with open(file_path, 'r', encoding='utf-8') as f_in, open(unk_path, 'w', encoding='utf-8') as f_out: + for line in f_in: + raw_tokens = line.strip().split(' ') + tokens = [] + for token in raw_tokens: + tokens.append(replace_number(token)) + # Starting each line with a blank line is required + # Some systems replace \n with and assume, like in PTB, everything is space separated + tokens = [''] + tokens + ['\n'] + line = ' '.join(tokens) + f_out.write(line) + + +def main(args): + + input_path = Path(args.input) + assert input_path.exists(), f'Error: {input_path} does not exist.' + + sml_wiki = input_path / f'{args.lang}-2' + lrg_wiki = input_path / f'{args.lang}-100' + assert sml_wiki.exists(), f'Error: {sml_wiki} does not exist.' + assert lrg_wiki.exists(), f'Error: {lrg_wiki} does not exist.' + + splits = ['train', 'valid', 'test'] + for wiki in [sml_wiki, lrg_wiki]: + for split in splits: + # replace numbers with placeholders + file_path = wiki / f'{args.lang}.wiki.{split}.tokens' + unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk' + replace_numbers(file_path, unk_path) + + sml_wiki_train = sml_wiki / f'{args.lang}.wiki.train.tokens' + lrg_wiki_train = lrg_wiki / f'{args.lang}.wiki.train.tokens' + + sml_vocab = build_vocab(sml_wiki_train) + print(f'{args.lang}-2 vocab size: {len(sml_vocab)}') + lrg_vocab = build_vocab(lrg_wiki_train) + print(f'{args.lang}-100 vocab size: {len(lrg_vocab)}') + + # replace words not in the vocab with + for wiki, vocab in zip([sml_wiki, lrg_wiki], [sml_vocab, lrg_vocab]): + for split in splits: + unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk' + limit_vocab(unk_path, vocab) + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser() + parser.add_argument('-i', '--input', required=True, + help='the directory of the wikitext files') + parser.add_argument('-l', '--lang', required=True, + help='the iso code of the language of the Wikipedia ' + 'documents, e.g. en, fr, de, etc.') + args = parser.parse_args() + main(args) diff --git a/ulmfit/prepare_wiki.sh b/ulmfit/prepare_wiki.sh new file mode 100644 index 0000000..c2fc481 --- /dev/null +++ b/ulmfit/prepare_wiki.sh @@ -0,0 +1,52 @@ +#!/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}/wiki_dumps" +EXTR_DIR="${ROOT}/wiki_extr" +WIKI_DIR="${ROOT}/wiki" +EXTR="wikiextractor" +mkdir -p "${ROOT}" +mkdir -p "${DUMP_DIR}" +mkdir -p "${EXTR_DIR}" +mkdir -p "${WIKI_DIR}" + +echo "Saving data in ""$ROOT" +read -r -p "Choose a language (e.g. en, bh, fr, etc.): " choice +LANG="$choice" +echo "Chosen language: ""$LANG" +DUMP_FILE="${LANG}wiki-latest-pages-articles.xml.bz2" +DUMP_PATH="${DUMP_DIR}/${DUMP_FILE}" + +if [ ! -f "${DUMP_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://dumps.wikimedia.org/""${LANG}""wiki/latest/""${DUMP_FILE}""" -P "${DUMP_DIR}" +else + echo "${DUMP_PATH} already exists. Skipping download." +fi + +# Check if directory exists +if [ ! -d "${EXTR}" ]; then + git clone https://github.com/attardi/wikiextractor.git + cd "${EXTR}" + python setup.py install +fi + +EXTR_PATH="${EXTR_DIR}/${LANG}" +if [ ! -d "${EXTR_PATH}" ]; then + read -r -p "Continue to extract Wikipedia (WARNING: This might take a long time!) (y/n)? " choice + case "$choice" in + y|Y ) echo "Extracting ${DUMP_PATH} to ${EXTR_PATH}...";; + n|N ) echo "Exiting";exit 1;; + * ) echo "Invalid answer";exit 1;; + esac + python wikiextractor/WikiExtractor.py -s --json -o "${EXTR_PATH}" "${DUMP_PATH}" +else + echo "${EXTR_PATH} already exists. Skipping extraction." +fi diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py new file mode 100644 index 0000000..157c0d5 --- /dev/null +++ b/ulmfit/pretrain_lm.py @@ -0,0 +1,134 @@ +""" +Script to train a model on a preprocessed Wiki dataset. Note that the dataset is +expected to have been tokenized with Moses and processed with `postprocess_wikitext.py`. +That is, the data is expected to be white-space separated and numbers are expected +to be split. +""" +import fire +import numpy as np + +from fastai import DataBunch, partial, optim, fit_one_cycle +from fastai.text import LanguageModelLoader, get_language_model, RNNLearner, TextLMDataBunch +import torch +from ulmfit.utils import read_file, read_whitespace_file,\ + DataStump, validate, PAD, UNK + +import pickle + +from pathlib import Path + +from collections import Counter + +# to install, do: +# conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92] +# cupy needs to be installed for QRNN + + +def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000, + bs=70, bptt=70, name='wt-103', model_dir='models', num_epochs=10): + """ + :param dir_path: The path to the directory of the file. + :param cuda_id: The id of the GPU. Uses GPU 0 by default or no GPU when + run on CPU. + :param qrrn: Use a QRNN. Requires installing cupy. + :param clean: Train on the clean + :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 model_dir: The path to the directory where the models should be saved + """ + if not torch.cuda.is_available(): + print('CUDA not available. Setting device=-1.') + cuda_id = -1 + torch.cuda.set_device(cuda_id) + + dir_path = Path(dir_path) + assert dir_path.exists() + model_dir = Path(model_dir) + model_dir.mkdir(exist_ok=True) + print('Batch size:', bs) + print('Max vocab:', max_vocab) + model_name = 'qrnn' if qrnn else 'lstm' + if qrnn: + print('Using QRNNs...') + + trn_path = dir_path / 'wiki.train.tokens' + val_path = dir_path / 'wiki.valid.tokens' + tst_path = dir_path / 'wiki.test.tokens' + for path_ in [trn_path, val_path, tst_path]: + assert path_.exists(), f'Error: {path_} does not exist.' + + if clean: + # read the already whitespace separated data without any preprocessing + trn_tok = read_whitespace_file(trn_path) + val_tok = read_whitespace_file(val_path) + + # create the vocabulary + cnt = Counter(word for sent in trn_tok for word in sent) + itos = [o for o,c in cnt.most_common(n=max_vocab)] + itos.insert(1, PAD) #  set pad id to 1 to conform to fast.ai standard + assert UNK in itos, f'Unknown words are expected to have been replaced with {UNK} in the data.' + stoi = {w: i for i, w in enumerate(itos)} + + 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]) + + # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos)) + trn_dl = LanguageModelLoader(DataStump(trn_ids), bs, bptt) + val_dl = LanguageModelLoader(DataStump(val_ids), bs, bptt) + data_lm = DataBunch(trn_dl, val_dl) + else: + # apply fastai preprocessing and tokenization + read_file(trn_path, 'train') + read_file(val_path, 'valid') + + data_lm = TextLMDataBunch.from_csv(dir_path, max_vocab=max_vocab) + itos = data_lm.train_ds.vocab.itos + stoi = data_lm.train_ds.vocab.stoi + + print('Size of vocabulary:', len(itos)) + print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)])) + + # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103) + # for training on smaller datasets, more dropout is necessary + if 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]) * 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]) * 0.1 + + model = get_language_model(len(itos), emb_sz, nh, nl, pad_token=1, input_p=dps[0], + output_p=dps[1], weight_p=dps[2], + embed_p=dps[3], hidden_p=dps[4], qrnn=qrnn) + learn = RNNLearner(data_lm, model, bptt, path=model_dir.parent, model_dir=model_dir.name, + clip=0.12) + + # save vocabulary + print('Saving vocabulary...') + with open(model_dir / f'itos_{name}.pkl', 'wb') as f: + pickle.dump(itos, f) + + # 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 + + fit_one_cycle(learn, num_epochs, 5e-3, (0.8, 0.7), wd=1e-7) + + if clean and max_vocab is None: + # only if we use the unpreprocessed version and the full vocabulary + # are the perplexity results comparable to previous work + 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(model, tst_ids, bptt) + print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item()) + + learn.save(f'{model_name}_{name}') + torch.save(learn.opt.opt.state_dict(), learn.path / learn.model_dir / f'{model_name}3_{name}_state.pth') + + +if __name__ == '__main__': + fire.Fire(pretrain_lm) diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py new file mode 100644 index 0000000..ddbbd60 --- /dev/null +++ b/ulmfit/train_clas.py @@ -0,0 +1,116 @@ +""" +Train a classifier on top of a language model trained with `pretrain_lm.py`. +Optionally fine-tune LM before. +""" +import numpy as np +import pickle + +from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner +from fastai import fit_one_cycle +from ulmfit.utils import PAD, UNK, read_imdb, PAD_TOKEN_ID + +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) + model_dir = Path(model_dir) + assert dir_path.exists(), f'Error: {dir_path} does not exist.' + assert model_dir.exists(), f'Error: {model_dir} does not exist.' + + if qrnn: + print('Using QRNNs...') + + 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)) + + # 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)) + + # 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) + + 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') + 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) + + # todo implemend save and load + #data_lm.save() + #data_clas.save() + #data_lm = TextLMDataBunch.load(path) + #data_clas = TextClasDataBunch.load(path, bs=bs) + + if qrnn: + emb_sz, nh, nl = 400, 1550, 3 + else: + emb_sz, nh, nl = 400, 1150, 3 + learn = language_model_learner( + data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, qrnn=qrnn, + pad_token=PAD_TOKEN_ID, + pretrained_fnames=(f'lstm_{pretrain_name}', f'itos_{pretrain_name}'), + path=model_dir.parent, model_dir=model_dir.name) + + if fine_tune: + print('Fine-tuning the language model...') + learn.unfreeze() + learn.fit(2, slice(1e-4, 1e-2)) + + # save encoder + learn.save_encoder('enc') + + learn = text_classifier_learner(data_clas, bptt=bptt, pad_token=PAD_TOKEN_ID, + path=model_dir.parent, model_dir=model_dir.name, + qrnn=qrnn, emb_sz=emb_sz, nh=nh, nl=nl) + + learn.load_encoder('enc') + fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7) + + learn.freeze_to(-2) + fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7) + + learn.unfreeze() + fit_one_cycle(learn, 10, 5e-3, (0.8, 0.7), wd=1e-7) + + +if __name__ == '__main__': fire.Fire(new_train_clas) diff --git a/ulmfit/utils.py b/ulmfit/utils.py new file mode 100644 index 0000000..5f40940 --- /dev/null +++ b/ulmfit/utils.py @@ -0,0 +1,127 @@ +""" +Utility methods for data processing. +""" +import pandas as pd +import numpy as np +from fastai import F, to_device +import torch +from tqdm import tqdm +import re +import csv + +EOS = '' +UNK = '' +PAD = '' +PAD_TOKEN_ID = 1 + +number_match_re = re.compile(r'^([0-9]+[,.]?)+$') +number_split_re = re.compile(r'([,.])') + + +def replace_number(token): + """Replaces a number and returns a list of one or multiple tokens.""" + if number_match_re.match(token): + return number_split_re.sub(r' @\1@ ', token) + return token + + +def read_file(file_path, outname): + """Reads a text file and writes it to a .csv.""" + with open(file_path, encoding='utf8') as f: + text = f.readlines() + df = pd.DataFrame( + {'text': np.array(text), 'labels': np.zeros(len(text))}, + columns=['labels', 'text']) + df.to_csv(file_path.parent / f'{outname}.csv', header=False, index=False) + + +def read_whitespace_file(filepath): + """Reads a file and prepares the tokens.""" + tokens = [] + with open(filepath, encoding='utf-8') as f: + for line in f: + # newlines are replaced with EOS + tokens.append(line.split() + [EOS]) + 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: + """Placeholder class as LanguageModelLoader requires object with ids attribute.""" + def __init__(self, ids): + self.ids = ids + self.loss_func = F.cross_entropy + + +def validate(model, ids, bptt=2000): + """ + Return the validation loss and perplexity of a model + :param model: model to test + :param ids: data on which to evaluate the model + :param bptt: bptt for this evaluation (doesn't change the result, only the speed) + From https://github.com/sgugger/Adam-experiments/blob/master/lm_val_fns.py#L34 + """ + data = TextReader(np.concatenate(ids), bptt) + model.eval() + model.reset() + total_loss, num_examples = 0., 0 + for inputs, targets in tqdm(data): + outputs, raws, outs = model(to_device(inputs, None)) + p_vocab = F.softmax(outputs, 1) + for i, pv in enumerate(p_vocab): + targ_pred = pv[targets[i]] + total_loss -= torch.log(targ_pred.detach()) + num_examples += len(inputs) + mean = total_loss / num_examples # divide by total number of tokens + return mean, np.exp(mean) + + +class TextReader(): + """ Returns a language model iterator that iterates through batches that are of length N(bptt,5) + The first batch returned is always bptt+25; the max possible width. This is done because of they way that pytorch + allocates cuda memory in order to prevent multiple buffers from being created as the batch width grows. + From: https://github.com/sgugger/Adam-experiments/blob/master/lm_val_fns.py#L3 + """ + def __init__(self, nums, bptt, backwards=False): + self.bptt,self.backwards = bptt,backwards + self.data = self.batchify(nums) + self.i,self.iter = 0,0 + self.n = len(self.data) + + def __iter__(self): + self.i,self.iter = 0,0 + while self.i < self.n-1 and self.iter