Initial import from fastaiv1

initial import from https://github.com/fastai/fastai/tree/ulmfit_v1/courses/dl2/imdb_scripts

.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

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.

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)

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 <UNK> 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(' ') + ['<eos>']
+            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 <eos> 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 <unk>
+    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)

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

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)

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)

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 = '<eos>'
+UNK = '<unk>'
+PAD = '<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<len(self):
+            res = self.get_batch(self.i, self.bptt)
+            self.i += self.bptt
+            self.iter += 1
+            yield res
+
+    def __len__(self): return self.n // self.bptt
+
+    def batchify(self, data):
+        data = np.array(data)[:,None]
+        if self.backwards: data=data[::-1]
+        return torch.LongTensor(data)
+
+    def get_batch(self, i, seq_len):
+        source = self.data
+        seq_len = min(seq_len, len(source) - 1 - i)
+        return source[i:i+seq_len], source[i+1:i+1+seq_len].view(-1)