diff --git a/README.md b/README.md
index 95bddb7..6b65884 100644
--- a/README.md
+++ b/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  
@@ -8,11 +46,12 @@ We have a fork of fastai to propose changes to fastai.text, with a branch for th
 Let us know that you want to start collaboration on fastai forum thread: [Multilingual ULMFIT](https://forums.fast.ai/t/multilingual-ulmfit/28117)
 and you will get access to both repositories.
 
-Follow the developer [installation of fastai] (https://github.com/fastai/fastai#developer-install)
-Add n-waves/fastai as additional remote as described here: https://help.github.com/articles/adding-a-remote/
-Here is what I've did:
-```bash
+- Follow the [developer installation of fastai](https://github.com/fastai/fastai#developer-install)
+- Add n-waves/fastai as additional remote as described here: https://help.github.com/articles/adding-a-remote/
 
+Here is what I did:
+```bash
+$ cd fastai
 $ git remote add n-waves https://github.com/n-waves/fastai.git
 $ git remote -v 
 n-waves	https://github.com/n-waves/fastai.git (fetch)
diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py
index 5f40940..6397ca5 100644
--- a/fastai_contrib/utils.py
+++ b/fastai_contrib/utils.py
@@ -3,20 +3,228 @@ 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
+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'
+}
 
 number_match_re = re.compile(r'^([0-9]+[,.]?)+$')
 number_split_re = re.compile(r'([,.])')
 
+CLASSES = ['neg', 'pos', 'unsup']
+
+
+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, message="One or more required files cannot be found."):
+    error = False
+    for path in paths:
+        if not path.exists():
+            print(f'Error: {path} does not exist.')
+            error = True
+    if error:
+        raise FileNotFoundError(message)
+
+def get_data_folder():
+    """
+    return data folder to use for future processing
+    """
+    return (pathlib.Path(__file__).parent.parent / "data")
+
+def get_scripts_folder():
+    """
+    return data folder to use for future processing
+    """
+    return (pathlib.Path(__file__).parent.parent)
+
+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."""
@@ -45,23 +253,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:
@@ -124,4 +315,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)
\ No newline at end of file
+        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
\ No newline at end of file
diff --git a/prepare_imdb.sh b/prepare_imdb.sh
new file mode 100644
index 0000000..6d2488c
--- /dev/null
+++ b/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_lm==False
+
diff --git a/prepare_wiki-en.sh b/prepare_wiki-en.sh
new file mode 100755
index 0000000..9e5add9
--- /dev/null
+++ b/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"
\ No newline at end of file
diff --git a/ulmfit/prepare_wiki.sh b/prepare_wiki.sh
similarity index 65%
rename from ulmfit/prepare_wiki.sh
rename to prepare_wiki.sh
index c2fc481..a43a868 100644
--- a/ulmfit/prepare_wiki.sh
+++ b/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
diff --git a/prepare_xnli.sh b/prepare_xnli.sh
new file mode 100644
index 0000000..8a04834
--- /dev/null
+++ b/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"
diff --git a/requirements.txt b/requirements.txt
index bb5bd5b..eb003e9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,4 @@
 fire>=0.1.3
-cupy>=5.0.0
\ No newline at end of file
+cupy>=5.0.0
+scikit-learn>=0.20
+sacremoses>=0.0.5
\ No newline at end of file
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
new file mode 100644
index 0000000..77c77be
--- /dev/null
+++ b/tests/test_end_to_end.py
@@ -0,0 +1,47 @@
+import fire
+import pytest
+
+import ulmfit.pretrain_lm
+import ulmfit.train_clas
+from fastai import *
+from fastai.text import *
+from fastai_contrib.utils import  *
+"""
+It is a mixture of a pytest unit test and woven together to compose an end to end functional test. 
+"""
+
+def check_data_exists():
+    data = get_data_folder()
+
+    wt2 = data / "wiki" / "wikitext-2"
+    imdb = data / "imdb"
+    ensure_paths_exists(wt2 / "en.wiki.train.tokens",
+                        imdb / "train.csv",
+                        message="We don't run data preparation scripts automatically as it takes ages, run prepare_wiki-en.sh & prepare_imdb.sh")
+    return imdb, wt2
+
+def test_pretrain_lm():
+    imdb,wt2 = check_data_exists()
+    lm_name="end-to-end-test-quick"
+    results = ulmfit.pretrain_lm.pretrain_lm(
+        dir_path=wt2,
+        lang='en',
+        qrnn=True,
+        clean=True,
+        max_vocab=1000,
+        bs=80,
+        num_epochs=1,
+        name=lm_name,
+        ds_pct=0.03
+    )
+    assert results['accuracy'] > 0.30
+
+    results = ulmfit.train_clas.new_train_clas(
+                data_dir=get_data_folder(),
+                lang='en', pretrain_name=lm_name, model_dir=wt2/'models',
+                qrnn=True,
+                fine_tune=True,
+                max_vocab=1000,
+                bs=20, bptt=70, name=lm_name+'-imdb-clas',
+                dataset='imdb',
+                ds_pct=0.03)
diff --git a/ulmfit/README.md b/ulmfit/README.md
index c7d60f8..355965d 100644
--- a/ulmfit/README.md
+++ b/ulmfit/README.md
@@ -1,2 +1,51 @@
-# todo 
-- [] create new docs
\ No newline at end of file
+# Todo 
+- [ ] Update these docs
+
+Getting Started
+---
+
+## Download and Extract the Wikipedia corpus
+
+In Linux, you can do all the following steps automatically with [prepare_wiki.sh](./prepare_wiki.sh)
+
+**Manual Instructions**
+
+We use the [WikiExtractor.py](http://medialab.di.unipi.it/wiki/Wikipedia_Extractor). It is a Python script that extracts and cleans text from a [Wikipedia database dump](http://download.wikimedia.org/).
+
+At the end of this step, you should have the following directory structure inside ulmfit:
+```bash
+
+|- data
+  |- wiki
+  |- wiki_dumps
+  |- wiki_extr
+|- wikiextractor
+```
+The extracted data should be in the folder `wiki_extr` -> language name e.g.`en` (english), `fr` (french) `hi` (hindi) and so on. 
+
+## Create and Post Process WikiText
+
+### Get the Dependencies
+
+**Python Fire**: To install Python Fire with pip, run: `pip install fire`
+
+To install Python Fire with conda, run: `conda install fire -c conda-forge`
+
+**Moses Tokenizer**: To install Moses Tokenizer: `pip install -U sacremoses`
+
+Use the Python script [create_wikitext.py](./create_wikitext.py) to process the extracted Wikipedia documents. 
+
+### 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/wiki/hi -l hi
+``` 
+for hindi (unicode: 'hi')
+
+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 data/wiki/hi-2 hi
+python postprocess_wikitext.py data/wiki/hi-100 hi
+```
diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py
index 5d24284..cca6a35 100644
--- a/ulmfit/create_wikitext.py
+++ b/ulmfit/create_wikitext.py
@@ -57,7 +57,7 @@ def write_wikitext(file_path, text_iter, mt, 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))
@@ -76,8 +76,10 @@ def main(args):
 
     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)
 
@@ -87,15 +89,18 @@ def main(args):
         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'
-
+        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)
+        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, 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__':
 
diff --git a/ulmfit/postprocess_wikitext.py b/ulmfit/postprocess_wikitext.py
index f609f9a..7ae8774 100644
--- a/ulmfit/postprocess_wikitext.py
+++ b/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)
\ No newline at end of file
diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py
index 55a64db..eba27c4 100644
--- a/ulmfit/pretrain_lm.py
+++ b/ulmfit/pretrain_lm.py
@@ -32,11 +32,11 @@ def accuracy_bwd(input, targs):
     return accuracy(input[...,1], targs[...,1])
 
 
-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,
-                bidir=False, ds_pct=1.0):
+def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
+                bs=70, bptt=70, name='wt-103', num_epochs=10,  bidir=False, ds_pct=1.0):
     """
-    :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.
     :param qrrn: Use a QRNN. Requires installing cupy.
@@ -48,6 +48,8 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
     :param model_dir: The path to the directory where the models should be saved
     :param bidir: whether the language model is bidirectional
     """
+    results = {}
+    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
@@ -63,9 +65,9 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
     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'
+    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.'
 
@@ -74,8 +76,9 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
         trn_tok = read_whitespace_file(trn_path)
         val_tok = read_whitespace_file(val_path)
         if ds_pct < 1.0:
-            trn_tok = trn_tok[:int(len(trn_tok) * ds_pct)]
-            val_tok = val_tok[:int(len(val_tok) * ds_pct)]
+            trn_tok = trn_tok[:max(20, int(len(trn_tok) * ds_pct))]
+            val_tok = val_tok[:max(20, int(len(val_tok) * ds_pct))]
+            print(f"Limiting data sets to {ds_pct*100}%, trn {len(trn_tok)}, val: {len(val_tok)}")
 
         # create the vocabulary
         cnt = Counter(word for sent in trn_tok for word in sent)
@@ -113,7 +116,7 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
     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
+        dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15])
         drop_mult = 0.1
     else:
         emb_sz, nh, nl = 400, 1150, 3
@@ -121,11 +124,11 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
         dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15])
         drop_mult = 0.1
 
-    fastai.text.learner.default_dropout['language'] = dps
+    fastai.text.learner.default_dropout['language'] = dps * drop_mult
 
     lm_learner = bilm_learner if bidir else language_model_learner
     learn = lm_learner(data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, pad_token=1,
-                       drop_mult=drop_mult, tie_weights=True,
+                       drop_mult=drop_mult, tie_weights=True, model_dir=model_dir,
                        bias=True, qrnn=qrnn, clip=0.12)
     # compared to standard Adam, we set beta_1 to 0.8
     learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99))
@@ -136,23 +139,34 @@ def pretrain_lm(dir_path, cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
     else:
         learn.metrics = [accuracy]
     # save vocabulary
-    print('Saving vocabulary...')
-    with open(model_dir / f'itos_{name}.pkl', 'wb') as f:
+    print(f"Saving vocabulary as {dir_path / model_dir}")
+    results['itos_fname'] = dir_path / model_dir / f'itos_{name}.pkl'
+    with open(results['itos_fname'], 'wb') as f:
         pickle.dump(itos, f)
 
     learn.fit_one_cycle(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
+
+        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])
-        logloss, perplexity = validate(model, tst_ids, bptt)
+        logloss, perplexity = validate(learn.model, tst_ids, bptt)
         print('Test logloss:', logloss.item(), 'perplexity:', perplexity.item())
 
+    print(f"Saving models at {learn.path / learn.model_dir}")
     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')
 
+    opt_state_path = learn.path / learn.model_dir / f'{model_name}3_{name}_state.pth'
+    print(f"Saving optimiser state at {opt_state_path}")
+    torch.save(learn.opt.opt.state_dict(), opt_state_path)
+
+    results['accuracy'] = learn.validate()[1]
+    return results
 
 if __name__ == '__main__':
     fire.Fire(pretrain_lm)
diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py
index caa71be..4bfbb8e 100644
--- a/ulmfit/train_clas.py
+++ b/ulmfit/train_clas.py
@@ -5,81 +5,118 @@ 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.
+    """
+    results={}
+    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'
+    lm_name = f'{model_name}_{pretrain_name}'
+    pretrained_fname = (lm_name, f'itos_{pretrain_name}')
 
-    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))
+    ensure_paths_exists(data_dir,
+                        dataset_dir,
+                        model_dir,
+                        model_dir/f"{pretrained_fname[0]}.pth",
+                        model_dir/f"{pretrained_fname[1]}.pkl")
 
-            # split off validation set if it does not exist
-            val_path = dir_path / 'valid.csv'
-            if not val_path.exists():
-                trn_len = int(len(trn_toks) * 0.9)
-                trn_toks, val_toks = trn_toks[:trn_len], trn_toks[trn_len:]
-                trn_lbls, val_lbls = trn_lbls[:trn_len], trn_lbls[trn_len:]
-            else:
-                val_toks, val_lbls = read_imdb(val_path, MosesTokenizer(lang))
 
-            # 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)
+    tmp_dir = dataset_dir / 'tmp'
+    tmp_dir.mkdir(exist_ok=True)
+    vocab_file = tmp_dir / f'vocab_{lang}.pkl'
 
-            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])
+    if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists():
+        print('Reading the data...')
+        toks, lbls = read_clas_data(dataset_dir, dataset, lang)
 
-            print(f'Train size: {len(trn_ids)}. Valid size: {len(val_ids)}. '
-                  f'Test size: {len(tst_ids)}.')
+        # 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)
 
-            for split, ids, lbl in zip(['train', 'valid', 'test'],
-                             [trn_ids, val_ids, tst_ids],
-                             [trn_lbls, val_lbls, tst_lbls]):
-                np.save(tmp_dir / f'{split}_ids.npy', ids)
-                np.save(tmp_dir / f'{split}_lbl.npy', lbl)
-
-        data_lm = TextLMDataBunch.from_id_files(tmp_dir, test='test')
-        data_clas = TextClasDataBunch.from_id_files(tmp_dir, test='test')
+        ids = {}
+        for split in [TRN, VAL, TST]:
+            ids[split] = np.array([([stoi.get(w, stoi[UNK]) for w in s])
+                                   for s in toks[split]])
+            np.save(tmp_dir / f'{split}_{lang}_ids.npy', ids[split])
+            np.save(tmp_dir / f'{split}_{lang}_lbl.npy', lbls[split])
     else:
-        # use fastai peprocessing and tokenization
-        data_lm = TextLMDataBunch.from_csv(dir_path, bs=bs)
-        data_clas = TextClasDataBunch.from_csv(dir_path, vocab=data_lm.train_ds.vocab, bs=bs)
+        print('Loading the pickled data...')
+        ids, lbls = {}, {}
+        for split in [TRN, VAL, TST]:
+            ids[split] = np.load(tmp_dir / f'{split}_{lang}_ids.npy')
+            lbls[split] = np.load(tmp_dir / f'{split}_{lang}_lbl.npy')
+        with open(vocab_file, 'rb') as f:
+            vocab = pickle.load(f)
 
-        # todo implemend save and load
-        #data_lm.save()
-        #data_clas.save()
-        #data_lm = TextLMDataBunch.load(path)
-        #data_clas = TextClasDataBunch.load(path, bs=bs)
+    print(f'Train size: {len(ids[TRN])}. Valid size: {len(ids[VAL])}. '
+          f'Test size: {len(ids[TST])}.')
+
+    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,29 +125,38 @@ 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:
+    lm_enc_finetuned  = f"{lm_name}_{dataset}_enc"
+    if fine_tune and not (model_dir / f"lm_enc_finetuned.pth").exists():
         print('Fine-tuning the language model...')
         learn.unfreeze()
         learn.fit(2, slice(1e-4, 1e-2))
 
         # save encoder
-        learn.save_encoder('enc')
+        learn.save_encoder(lm_enc_finetuned)
 
+    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')
-    fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7)
+    learn.load_encoder(lm_enc_finetuned)
+
+    learn.fit_one_cycle(1, 2e-2, moms=(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.fit_one_cycle(1, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
+
+    learn.freeze_to(-3)
+    learn.fit_one_cycle(1, slice(5e-3 / (2.6 ** 4), 5e-3), moms=(0.8, 0.7), wd=1e-7)
 
     learn.unfreeze()
-    fit_one_cycle(learn, 10, 5e-3, (0.8, 0.7), wd=1e-7)
+    learn.fit_one_cycle(2, slice(1e-3 / (2.6 ** 4), 1e-3), moms=(0.8, 0.7), wd=1e-7)
+    results['accuracy'] = learn.validate()[1]
+    print(f"Saving models at {learn.path / learn.model_dir}")
+    learn.save(f'{model_name}_{name}')
+    return results
 
-
-if __name__ == '__main__': fire.Fire(new_train_clas)
+if __name__ == '__main__':
+    fire.Fire(new_train_clas)