diff --git a/ulmfit/prepare_wiki.sh b/prepare_wiki.sh
similarity index 61%
rename from ulmfit/prepare_wiki.sh
rename to prepare_wiki.sh
index c2fc481..30ca7eb 100644
--- a/ulmfit/prepare_wiki.sh
+++ b/prepare_wiki.sh
@@ -3,8 +3,18 @@
 
 # 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"
+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"
 EXTR="wikiextractor"
 mkdir -p "${ROOT}"
@@ -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/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 a4b2f4e..6370955 100644
--- a/ulmfit/pretrain_lm.py
+++ b/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])