Merge branch 'master' into sentencepiece

README.md

@@ -1,6 +1,44 @@
 # ulmfit-multilingual
 Temporary repository used for collaboration on application of for multiple languages.
 
+## data directory strucutre
+
+Directory structure after changes to the way we process wiki dumps.
+```
+data
+├── imdb
+│   ├── aclImdb
+│   ├── imdb_lm
+│   └── tmp
+├── wiki
+│   ├── de-100
+│   │   └── models
+│   ├── de-100-unk
+│   │   └── models
+│   ├── de-2
+│   │   └── models
+│   ├── de-2-unk
+│   │   └── models
+│   ├── de-all
+│   │   └── models
+│   ├── wikitext-103
+│   │   └── models
+│   └── wikitext-2
+│       └── models
+├── wiki_dumps
+├── wiki_extr
+│   └── de
+│       ├── AA
+│       ├── AB
+...
+        └── CC
+└── xnli
+    ├── XNLI-1.0
+    └── XNLI-MT-1.0
+        ├── multinli
+        └── xnli
+```
+
 ## how to contribute
 We have a fork of fastai to propose changes to fastai.text, with a branch for this project:
  https://github.com/n-waves/fastai/tree/ulmfit_multilingual  

fastai_contrib/utils.py

@@ -3,21 +3,39 @@ Utility methods for data processing.
 """
 import pandas as pd
 import numpy as np
+import fire
 from fastai import F, to_device
 import torch
 from tqdm import tqdm
 import re
 import csv
+
 from functools import reduce
 from fastai.text.data import TextDataset
 from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab, default_rules
 from fastai.torch_core import *
 from pathlib import Path
 
+import pathlib
+import tarfile
+from sklearn import model_selection
+from sacremoses import MosesTokenizer
+from typing import Dict, Tuple, List
+
 EOS = '<eos>'
 UNK = '<unk>'
 PAD = '<pad>'
+SEP = '<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'
+}
+
+CLASSES = ['neg', 'pos', 'unsup']
 
 number_match_re = re.compile(r'^([0-9]+[,.]?)+$')
 number_split_re = re.compile(r'([,.])')
@@ -74,6 +92,185 @@ def get_sentencepiece(path:PathOrStr, dataset:TextDataset, rules:ListRules=None,
     
     return {'tokenizer': tokenizer, 'vocab': vocab}
 
+def get_texts(path):
+    texts, labels = [],[]
+    for idx, label in enumerate(CLASSES):
+        for fname in (path/label).glob('*.*'):
+            texts.append(fname.open('r', encoding='utf-8').read())
+            labels.append(idx)
+    return np.array(texts), np.array(labels)
+
+
+def ensure_paths_exists(*paths):
+    error = False
+    for path in paths:
+        if not path.exists():
+            print(f'Error: {path} does not exist.')
+            error = True
+    if error:
+        raise FileNotFoundError("One or more required files cannot be found.")
+
+def prepare_imdb(file_path: str, prepare_lm = False):
+    """
+    function to extract aclImdb and combine into fastai standard format of labels and then text
+    columns
+
+    Args:
+        file_path: path to the aclImdb.tgz
+        prepare_lm (bool): prepare file for language model finetuning
+
+    Returns:
+        None
+    """
+
+    file_path = pathlib.Path(file_path)
+    dir_path = pathlib.Path(file_path.parent / 'aclImdb').resolve()
+    assert tarfile.is_tarfile(file_path), "this is not a valid targz file"
+
+    if not dir_path.exists():
+        print(f"Extracting {file_path} to {dir_path}. This may take a long time...")
+        tgz_file = tarfile.open(file_path)
+        tgz_file.extractall(path=dir_path.parent) # the aclImdb.tgz has aclImdb dir packed
+        assert dir_path.exists()
+        print(f"Extracted to {dir_path}")
+
+    CLAS_PATH = dir_path.parent
+    CLAS_PATH.mkdir(exist_ok=True)
+
+    LM_PATH = dir_path.parent /'imdb_lm'
+    LM_PATH.mkdir(exist_ok=True)
+
+    # processing the split files to create train.csv and test.csv in fastai format
+    col_names = ['labels', 'text']
+    trn_texts, trn_labels = get_texts(dir_path/ 'train')
+    val_texts, val_labels = get_texts(dir_path / 'test')
+    np.random.seed(42)
+    trn_idx = np.random.permutation(len(trn_texts))
+    val_idx = np.random.permutation(len(val_texts))
+    trn_texts = trn_texts[trn_idx]
+    val_texts = val_texts[val_idx]
+    trn_labels = trn_labels[trn_idx]
+    val_labels = val_labels[val_idx]
+
+    df_trn = pd.DataFrame({'text': trn_texts, 'labels': trn_labels}, columns=col_names)
+    df_val = pd.DataFrame({'text': val_texts, 'labels': val_labels}, columns=col_names)
+    print(f"df_trn has {len(df_trn)} rows, while df_val has {len(df_val)} rows")
+    print(f"Writing them to {CLAS_PATH}")
+    df_trn[df_trn['labels'] != 2].to_csv(CLAS_PATH / 'train.csv', header=False, index=False)
+    df_val.to_csv(CLAS_PATH / 'test.csv', header=False, index=False)
+
+    (CLAS_PATH / 'classes.txt').open('w', encoding='utf-8').writelines(f'{o}\n' for o in CLASSES)
+
+    if prepare_lm:
+        print("Preparing LM data")
+        trn_texts, val_texts = model_selection.train_test_split(
+            np.concatenate([trn_texts, val_texts]), test_size=0.1)
+        print(f"trn_texts has {len(trn_texts)} samples, while val_texts has {len(val_texts)} rows")
+        print(f"Writing them to {LM_PATH}")
+        df_trn = pd.DataFrame({'text': trn_texts, 'labels': [0] * len(trn_texts)}, columns=col_names)
+        df_val = pd.DataFrame({'text': val_texts, 'labels': [0] * len(val_texts)}, columns=col_names)
+
+        df_trn.to_csv(LM_PATH / 'train.csv', header=False, index=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):
@@ -101,23 +298,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:
@@ -180,4 +360,8 @@ class TextReader():
     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)
+        return source[i:i+seq_len], source[i+1:i+1+seq_len].view(-1)
+
+
+if __name__ == "__main__":
+    fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz

prepare_imdb.sh

@@ -0,0 +1,11 @@
+#!/usr/bin/env bash
+
+ROOT="data"
+DATA_DIR="${ROOT}/imdb"
+mkdir -p "${DATA_DIR}"
+echo "Saving data in $DATA_DIR"
+wget -c "http://files.fast.ai/data/aclImdb.tgz" -P "${DATA_DIR}"
+
+echo "Imdb is raw text so we are tokenizing it with Moses"
+python -m fastai_contrib.utils prepare_imdb "${DATA_DIR}/aclImdb.tgz"
+

prepare_wiki-en.sh

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+ROOT="data"
+DATA_DIR="${ROOT}/wiki/"
+mkdir -p "${DATA_DIR}"
+echo "Saving data in ""$DATA_DIR"
+wget -c "https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip" -P "${DATA_DIR}"
+wget -c "https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-v1.zip" -P  "${DATA_DIR}"
+
+unzip "${DATA_DIR}/wikitext-2-v1.zip" -d "${DATA_DIR}"
+unzip "${DATA_DIR}/wikitext-103-v1.zip" -d "${DATA_DIR}"
+
+for f in ${DATA_DIR}/wikitext-*/wiki.*.tokens; do
+    nf=$(dirname $f)/en.$(basename $f)
+    echo "Renaming $f to $nf"
+    mv $f $(dirname $f)/en.$(basename $f)
+done
+
+echo "Please note wikitext en is already tokenized with Moses"

prepare_wiki.sh

@@ -3,6 +3,16 @@
 
 # Script is partially based on https://github.com/facebookresearch/fastText/blob/master/get-wikimedia.sh
 ROOT="data"
+echo "Saving data in ""$ROOT"
+
+if [ "$1" == "" ] ; then
+    read -r -p "Choose a language (e.g. en, bh, fr, etc.): " choice
+    LANG="$choice"
+else
+    LANG="$1"
+fi
+echo "Chosen language: ""$LANG"
+
 DUMP_DIR="${ROOT}/wiki_dumps"
 EXTR_DIR="${ROOT}/wiki_extr"
 WIKI_DIR="${ROOT}/wiki"
@@ -12,20 +22,10 @@ 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."
@@ -36,17 +36,18 @@ if [ ! -d "${EXTR}" ]; then
   git clone https://github.com/attardi/wikiextractor.git
   cd "${EXTR}"
   python setup.py install
+  cd ..
 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
+
+python -m ulmfit.create_wikitext -i "${EXTR_PATH}"  -l "${LANG}" -o "${WIKI_DIR}"
+
+python -m ulmfit.postprocess_wikitext "${WIKI_DIR}/${LANG}-2" $LANG
+python -m ulmfit.postprocess_wikitext "${WIKI_DIR}/${LANG}-100" $LANG
+#python -m ulmfit.postprocess_wikitext "${WIKI_DIR}/${LANG}-all" $LANG

prepare_xnli.sh

@@ -0,0 +1,26 @@
+#!/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}"
+
+echo "Please note xnli en is already tokenized with Moses"

requirements.txt

@@ -1,2 +1,4 @@
 fire>=0.1.3
-cupy>=5.0.0
+cupy>=5.0.0
+scikit-learn>=0.20
+sacremoses>=0.0.5

ulmfit/README.md

@@ -38,7 +38,7 @@ Use the Python script [create_wikitext.py](./create_wikitext.py) to process the
 ### Create and Post-Process 
 If you used the automated shell script from previous step, this might look something like
 ```bash
-python create_wikitext.py -i data/wiki_extr/hi -o data/hindi -l hi
+python create_wikitext.py -i data/wiki_extr/hi -o data/wiki/hi -l hi
 ``` 
 for hindi (unicode: 'hi')
 
@@ -46,5 +46,6 @@ This should create two splits of your Wikimedia Dumps: a small and large one.
 
 _**Then**_, use the [postprocess_wikitext.py](./postprocess_wikitext.py) script to finish post processing. This processes numbers, builds a vocab, and limits the vocabulary size. This might look following for Hindi (`hi`)
 ```bash
-python postprocess_wikitext.py -i data/hindi -l hi
+python postprocess_wikitext.py data/wiki/hi-2 hi
+python postprocess_wikitext.py data/wiki/hi-100 hi
 ```

ulmfit/create_wikitext.py

@@ -6,14 +6,11 @@ Articles are tokenized using the Moses tokenizer. Articles with least than
 """
 import argparse
 from pathlib import Path
-from collections import Counter
 import json
 
 from shutil import copyfile
 
 from sacremoses import MosesTokenizer
-from fastai_contrib.utils import replace_number, UNK
-from fastai_contrib.tokenizers import get_sentencepiece, SentencepieceTokenizer
 
 
 def get_texts(root):
@@ -30,12 +27,10 @@ def get_texts(root):
                     yield text
 
 
-def write_wikitext(file_path, text_iter, tok, num_tokens, mode='w'):
-    
+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):
 
@@ -44,7 +39,7 @@ def write_wikitext(file_path, text_iter, tok, num_tokens, mode='w'):
             paragraphs = text.split('\n')
 
             for paragraph in paragraphs:
-                tokenized = tok.tokenize(paragraph.strip(), return_str=True)
+                tokenized = mt.tokenize(paragraph.strip(), return_str=True)
                 tokenized_paragraphs.append(tokenized)
 
                 tokens = tokenized.split(' ')  # split on whitespace to keep newlines
@@ -62,77 +57,13 @@ def write_wikitext(file_path, text_iter, tok, num_tokens, mode='w'):
                 f_out.write(tokenized + '\n')
 
             total_num_tokens += num_tokens_article + 1
-            if total_num_tokens > num_tokens:
+            if num_tokens is not None and 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 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(text_iter, 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. Writing to {unk_path}.')
-    with open(unk_path, 'w', encoding='utf-8') as f:
-        for text in text_iter:
-            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.write(line)
-
 
 def main(args):
 
@@ -141,52 +72,35 @@ def main(args):
     assert input_path.exists(), f'Error: {input_path} does not exist.'
     output.mkdir(exist_ok=True)
 
-    if args.subword:
-        #  TO DO load the text corpus
-        #  TO DO make get_sentencepiece return path to spm model
-        spm_path = get_sentencepiece(output, corpus)
-        tok = SentencepieceTokenizer(spm_path)
-    else:
-        tok = MosesTokenizer(args.lang)
+    mt = MosesTokenizer(args.lang)
 
     sml_wiki = output / f'{args.lang}-2'
     lrg_wiki = output / f'{args.lang}-100'
+    all_wiki = output / f'{args.lang}-all'
     sml_wiki.mkdir(exist_ok=True)
     lrg_wiki.mkdir(exist_ok=True)
+    all_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):
-        # TO DO maybe replace the numbers before tokenizing
-        unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk'
-        replace_numbers(text_iter, unk_path)
-
         sml_file_path = sml_wiki / f'{args.lang}.wiki.{split}.tokens'
-        write_wikitext(sml_file_path, text_iter, tok, token_num)
+        write_wikitext(sml_file_path, text_iter, mt, token_num)
         lrg_file_path = lrg_wiki / f'{args.lang}.wiki.{split}.tokens'
-
+        all_file_path = all_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}.')
+        print(f'Copying {sml_file_path} to {lrg_file_path} & {all_file_path}.')
         copyfile(sml_file_path, lrg_file_path)
-
-    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)}')
+        copyfile(sml_file_path, all_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, tok, 98000000, mode='a')
-
-    # replace words not in the vocab with <unk>
-    if not args.subword:
-        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)
-
+    write_wikitext(lrg_wiki_train, text_iter, mt, 98000000, mode='a')
+    all_wiki_train = all_wiki / f'{args.lang}.wiki.train.tokens'
+    copyfile(lrg_wiki_train, all_wiki_train)
+    write_wikitext(lrg_wiki_train, text_iter, mt,  None, mode='a') # TODO fix it (change lrg to all)
 
 if __name__ == '__main__':
 
@@ -201,8 +115,5 @@ if __name__ == '__main__':
     parser.add_argument('-l', '--lang', required=True,
                         help='the iso code of the language of the Wikipedia '
                              'documents, e.g. en, fr, de, etc.')
-    parser.add_argument('-sw', '--subword', default=False,
-                        help='set to use sub-word tokenization (sentencepiece)'
-                              'default tokenization method is Moses.')
     args = parser.parse_args()
     main(args)

ulmfit/postprocess_wikitext.py

@@ -7,6 +7,9 @@ import argparse
 
 from collections import Counter
 from pathlib import Path
+
+import fire
+
 from fastai_contrib.utils import replace_number, UNK
 
 
@@ -76,46 +79,27 @@ def replace_numbers(file_path, unk_path):
             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.'
+def postprocess_wikitext(path, lang):
 
+    wiki_path = Path(path)
+    assert wiki_path.exists(), f'Error: {wiki_path} does not exist.'
+    dest_path = wiki_path.parent / (wiki_path.name + "-unk")
+    dest_path.mkdir(exist_ok=True)
     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)}')
+    for split in splits:
+        # replace numbers with placeholders
+        file_path = wiki_path / f'{lang}.wiki.{split}.tokens'
+        assert file_path.exists(), f"Error: {file_path} does not exist."
+        unk_path = dest_path / file_path.name
+        replace_numbers(file_path, unk_path)
 
     # 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)
-
+    wiki_train = dest_path / f'{lang}.wiki.train.tokens'
+    vocab = build_vocab(wiki_train)
+    print(f'{wiki_path} vocab size: {len(vocab)}')
+    for split in splits:
+        unk_path = dest_path / f'{lang}.wiki.{split}.tokens'
+        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)
+    fire.Fire(postprocess_wikitext)

ulmfit/pretrain_lm.py

@@ -26,9 +26,9 @@ from collections import Counter
 
 
 def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
-                bs=70, bptt=70, name='wt-103', model_dir='models', num_epochs=10):
+                bs=70, bptt=70, name='wt-103', num_epochs=10):
     """
-    :param dir_path: The path to the directory of the file.
+    :param dir_path: The path to the directory that contains wiki text
     :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.
@@ -40,6 +40,8 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
     :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
     """
+
+    model_dir = 'models' # removed from params, as it is absolute models location in train_clas and here it is relative
     if not torch.cuda.is_available():
         print('CUDA not available. Setting device=-1.')
         cuda_id = -1
@@ -55,9 +57,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
     if qrnn:
         print('Using QRNNs...')
 
-    trn_path = dir_path / f'{lang}.wiki.train.tokens.unk'
-    val_path = dir_path / f'{lang}.wiki.valid.tokens.unk'
-    tst_path = dir_path / f'{lang}.wiki.test.tokens.unk'
+    trn_path = dir_path / f'{lang}.wiki.train.tokens'
+    val_path = dir_path / f'{lang}.wiki.valid.tokens'
+    tst_path = dir_path / f'{lang}.wiki.test.tokens'
     for path_ in [trn_path, val_path, tst_path]:
         assert path_.exists(), f'Error: {path_} does not exist.'
 
@@ -124,6 +126,7 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
     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
+
         print(f"Validating model performance with test tokens from: {trn_path}")
         tst_tok = read_whitespace_file(trn_path)
         tst_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in tst_tok])

ulmfit/train_clas.py

@@ -5,81 +5,117 @@ 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, ensure_paths_exists
+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='wt103', model_dir='models',
+                   qrnn=False,
+                   fine_tune=True, max_vocab=30000, bs=20, bptt=70, name='imdb-clas',
+                   dataset='imdb', ds_pct=1.0):
+    """
+    :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 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))
+    pretrained_fname = (f'{model_name}_{pretrain_name}', f'itos_{pretrain_name}')
 
-            # 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))
+    ensure_paths_exists(data_dir,
+                        dataset_dir,
+                        model_dir,
+                        model_dir/f"{pretrained_fname[0]}.pth",
+                        model_dir/f"{pretrained_fname[1]}.pkl")
 
-            # 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])
+    tmp_dir = dataset_dir / 'tmp'
+    tmp_dir.mkdir(exist_ok=True)
+    vocab_file = tmp_dir / f'vocab_{lang}.pkl'
 
-            print(f'Train size: {len(trn_ids)}. Valid size: {len(val_ids)}. '
-                  f'Test size: {len(tst_ids)}.')
+    if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists():
+        print('Reading the data...')
+        toks, lbls = read_clas_data(dataset_dir, dataset, lang)
 
-            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)
+        # 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)
 
-        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])}.')
+
+    if ds_pct < 1.0:
+        print(f"Makeing the dataset smaller {ds_pct}")
+        for split in [TRN, VAL, TST]:
+            ids[split] = ids[split][:int(len(ids[split])*ds_pct)]
+
+
+    data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=ids[TRN],
+                                       valid_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, train_ids=ids[TRN], valid_ids=ids[VAL],
+        train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs)
 
     if qrnn:
         emb_sz, nh, nl = 400, 1550, 3
@@ -88,10 +124,10 @@ def new_train_clas(dir_path, lang='en', pretrain_name='wt-103', model_dir='model
     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}'),
+        pretrained_fnames=pretrained_fname,
         path=model_dir.parent, model_dir=model_dir.name)
 
-    if fine_tune:
+    if fine_tune and not (model_dir / "enc.pth").exists():
         print('Fine-tuning the language model...')
         learn.unfreeze()
         learn.fit(2, slice(1e-4, 1e-2))
@@ -99,18 +135,27 @@ def new_train_clas(dir_path, lang='en', pretrain_name='wt-103', model_dir='model
         # save encoder
         learn.save_encoder('enc')
 
+    print("Starting classifier training")
     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')
+
+    torch.cuda.empty_cache()
     fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7)
 
+    torch.cuda.empty_cache()
     learn.freeze_to(-2)
     fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7)
 
+    torch.cuda.empty_cache()
     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)