Add end to end test and improve the train_clas params

fastai_contrib/utils.py

@@ -43,14 +43,26 @@ def get_texts(path):
     return np.array(texts), np.array(labels)
 
 
-def ensure_paths_exists(*paths):
+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("One or more required files cannot be found.")
+        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):
     """

prepare_imdb.sh

@@ -7,5 +7,5 @@ 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"
+python -m fastai_contrib.utils prepare_imdb "${DATA_DIR}/aclImdb.tgz" --prepare_lm==False
 

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)

ulmfit/pretrain_lm.py

@@ -26,7 +26,7 @@ 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', num_epochs=10):
+                bs=70, bptt=70, name='wt-103', num_epochs=10, ds_pct=1.0):
     """
     :param dir_path: The path to the directory that contains wiki text
     :param lang: the language unicode
@@ -40,7 +40,7 @@ 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
     """
-
+    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.')
@@ -67,6 +67,10 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
         # read the already whitespace separated data without any preprocessing
         trn_tok = read_whitespace_file(trn_path)
         val_tok = read_whitespace_file(val_path)
+        if ds_pct < 1.0:
+            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)
@@ -118,11 +122,14 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
 
     # save vocabulary
     print(f"Saving vocabulary as {dir_path / model_dir}")
-    with open(dir_path / model_dir / f'itos_{name}.pkl', 'wb') as f:
+    results['itos_fname'] = dir_path / model_dir / f'itos_{name}.pkl'
+    with open(results['itos_fname'], 'wb') as f:
         pickle.dump(itos, f)
 
     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
@@ -140,6 +147,8 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
     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)

ulmfit/train_clas.py

@@ -37,6 +37,7 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
                     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
@@ -57,8 +58,8 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
     if qrnn:
         print('Using QRNNs...')
     model_name = 'qrnn' if qrnn else 'lstm'
-
-    pretrained_fname = (f'{model_name}_{pretrain_name}', f'itos_{pretrain_name}')
+    lm_name = f'{model_name}_{pretrain_name}'
+    pretrained_fname = (lm_name, f'itos_{pretrain_name}')
 
     ensure_paths_exists(data_dir,
                         dataset_dir,
@@ -126,36 +127,36 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
         pad_token=PAD_TOKEN_ID,
         pretrained_fnames=pretrained_fname,
         path=model_dir.parent, model_dir=model_dir.name)
-
-    if fine_tune and not (model_dir / "enc.pth").exists():
+    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')
+    learn.load_encoder(lm_enc_finetuned)
 
-    torch.cuda.empty_cache()
-    fit_one_cycle(learn, 1, 5e-3, (0.8, 0.7), wd=1e-7)
+    learn.fit_one_cycle(1, 2e-2, moms=(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)
+    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)
 
-    torch.cuda.empty_cache()
     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)