Clean ups and fixes

fastai_contrib/data.py

@@ -5,8 +5,7 @@ from fastai.text.transform import *
 from fastai.basic_data import *
 from fastai.data_block import *
 
-
-###################### UPDATED CODE
+#region Modified fastai classes
 
 LanguageModelType=Enum('LanguageModelType', 'FwdLM BwdLM BiLM')
 
@@ -24,7 +23,7 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
         if getattr(self.dataset, 'item', None) is not None:
             yield LongTensor(getattr(self.dataset, 'item')).unsqueeze(1),LongTensor([0])
         idx = np.random.permutation(len(self.dataset)) if self.shuffle else range(len(self.dataset))
-        data = self.batchify(np.concatenate([self.dataset.x.items[i] for i in idx]))
+        data = self.batchify(np.concatenate([np.array(self.dataset.x.items[i], dtype=np.int) for i in idx]))
 
         pos, itr = 0,0
         while pos < self.n-1 and itr<len(self):
@@ -57,8 +56,10 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
         y = y.view(-1, 2) if self.lm_type == LanguageModelType.BiLM else y.view(-1)
         return x,y
 
-###################### NEW CODE
-
+#endregion
+#region Replaces fastai classes
 
 import fastai.text.data
 fastai.text.data.LanguageModelLoader = LanguageModelLoader # Replace original LanguageModelLoader with new verion
+
+#endregion

fastai_contrib/learner.py

@@ -5,6 +5,7 @@ from fastai.datasets import untar_data
 from fastai_contrib.models import get_bilm, get_rnn_classifier, get_birnn_classifier
 from fastai.text.learner import *
 
+#region New code
 
 def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:int=3, pad_token:int=1,
                   drop_mult:float=1., tie_weights:bool=True, bias:bool=True, qrnn:bool=False, pretrained_model=None,
@@ -26,11 +27,49 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in
         learn.freeze()
     return learn
 
+def bilm_text_classifier_learner(data: DataBunch, bptt: int = 70, max_len: int = 70 * 20, emb_sz: int = 400,
+                            nh: int = 1150, nl: int = 3,
+                            lin_ftrs: Collection[int] = None, ps: Collection[float] = None, pad_token: int = 1,
+                            drop_mult: float = 1., qrnn: bool = False, **kwargs) -> 'TextClassifierLearner':
+    "Create a RNN classifier."
+    dps = default_dropout['classifier'] * drop_mult
+    if lin_ftrs is None: lin_ftrs = [50]
+    if ps is None:  ps = [0.1]
+    ds = data.train_ds
+    vocab_size, n_class = len(data.vocab.itos), data.c
+    layers = [emb_sz * 3] + lin_ftrs + [n_class]
+    ps = [dps[4]] + ps
+    model = get_birnn_classifier(bptt, max_len, n_class, vocab_size, emb_sz, nh, nl, pad_token,
+                               layers, ps, input_p=dps[0], weight_p=dps[1], embed_p=dps[2], hidden_p=dps[3],
+                               qrnn=qrnn)
+    learn = RNNLearner(data, model, bptt, split_func=birnn_classifier_split, **kwargs)
+    return learn
 
+def bilm_split(model:nn.Module) -> List[nn.Module]:
+    "Split a RNN `model` in groups for differential learning rates."
+
+    return [f+b for f,b in zip(lm_split(model.fwd_lm),lm_split(model.bwd_lm))]
+
+def birnn_classifier_split(model:nn.Module) -> List[nn.Module]:
+    "Split a RNN `model` in groups for differential learning rates."
+    f_rnn,b_rnn = model[0].fwd_lm,model[0].bwd_lm
+    groups = [[f_rnn.encoder, f_rnn.encoder_dp,b_rnn.encoder, b_rnn.encoder_dp]]
+    groups += [a for a in zip(f_rnn.rnns, f_rnn.hidden_dps, b_rnn.rnns, b_rnn.hidden_dps, )]
+    groups.append([model[1]])
+    return groups
+
+def accuracy_fwd(input, targs):
+    return accuracy(input[...,0], targs[...,0])
+
+def accuracy_bwd(input, targs):
+    return accuracy(input[...,1], targs[...,1])
+
+
+#endregion
+#region Modified fastai code
 
 def convert_weights(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:Collection[str]) -> Weights:
     "Convert the model weights to go with a new vocabulary."
-    print(wgts.keys())
     if 'fwd_lm.0.encoder.weight' in wgts: #todo share embedding matrix computation
         wgts = convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='fwd_lm.')
         return convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='bwd_lm.')
@@ -53,49 +92,10 @@ def convert_weights_with_prefix(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:
     wgts[prefix+'1.decoder.bias'] = new_b
     return wgts
 
+#endregion
+#region Replace code in fastai
 
-
-def bilm_text_classifier_learner(data: DataBunch, bptt: int = 70, max_len: int = 70 * 20, emb_sz: int = 400,
-                            nh: int = 1150, nl: int = 3,
-                            lin_ftrs: Collection[int] = None, ps: Collection[float] = None, pad_token: int = 1,
-                            drop_mult: float = 1., qrnn: bool = False, **kwargs) -> 'TextClassifierLearner':
-    "Create a RNN classifier."
-    dps = default_dropout['classifier'] * drop_mult
-    if lin_ftrs is None: lin_ftrs = [50]
-    if ps is None:  ps = [0.1]
-    ds = data.train_ds
-    vocab_size, n_class = len(data.vocab.itos), data.c
-    layers = [emb_sz * 3] + lin_ftrs + [n_class]
-    ps = [dps[4]] + ps
-    model = get_birnn_classifier(bptt, max_len, n_class, vocab_size, emb_sz, nh, nl, pad_token,
-                               layers, ps, input_p=dps[0], weight_p=dps[1], embed_p=dps[2], hidden_p=dps[3],
-                               qrnn=qrnn)
-    learn = RNNLearner(data, model, bptt, split_func=birnn_classifier_split, **kwargs)
-    return learn
-
-
-def bilm_split(model:nn.Module) -> List[nn.Module]:
-    "Split a RNN `model` in groups for differential learning rates."
-
-    return [f+b for f,b in zip(lm_split(model.fwd_lm),lm_split(model.bwd_lm))]
-
-def birnn_classifier_split(model:nn.Module) -> List[nn.Module]:
-    "Split a RNN `model` in groups for differential learning rates."
-    f_rnn,b_rnn = model[0].fwd_lm,model[0].bwd_lm
-    groups = [[f_rnn.encoder, f_rnn.encoder_dp,b_rnn.encoder, b_rnn.encoder_dp]]
-    groups += [a for a in zip(f_rnn.rnns, f_rnn.hidden_dps, b_rnn.rnns, b_rnn.hidden_dps, )]
-    groups.append([model[1]])
-    return groups
-
-
-def accuracy_fwd(input, targs):
-    return accuracy(input[...,0], targs[...,0])
-
-
-def accuracy_bwd(input, targs):
-    return accuracy(input[...,1], targs[...,1])
-
-
-## Replace code in fastai
 import fastai.text.learner
-fastai.text.learner.convert_weights = convert_weights
+fastai.text.learner.convert_weights = convert_weights
+
+#endregion

fastai_contrib/models.py

@@ -2,6 +2,8 @@ from fastai.torch_core import *
 from fastai.layers import *
 from fastai.text.models import *
 
+#region New code
+
 class BiLMModel(nn.Module):
 
     def __init__(self, fwd_lm:nn.Module, bwd_lm:nn.Module):
@@ -145,4 +147,6 @@ def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_s
 
     model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), AvgPoolingLinearClassifier(layers, drops))
     model.reset()
-    return model
+    return model
+
+#endregion

fastai_contrib/utils.py

@@ -55,7 +55,7 @@ class SentencepieceTokenizer(BaseTokenizer):
         pass
 
 
-def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=None,
+def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, pre_rules:ListRules=None, post_rules:ListRules=None,
                       vocab_size:int=30000, model_type:str='unigram', input_sentence_size:int=1E7, 
                       pad_idx:int=PAD_TOKEN_ID):
     try:
@@ -67,15 +67,15 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N
     cache_name = 'tmp'
     os.makedirs(path / cache_name, exist_ok=True)
     os.makedirs(path / 'models', exist_ok=True)
-    rules = rules if rules is not None else []
-
+    pre_rules = pre_rules if pre_rules is not None else []
+    post_rules = post_rules if post_rules is not None else []
     
     # load the text frmo the train tokens file
     text = [line.rstrip('\n') for line in open(trn_path)]
     text = list(filter(None, text))
 
     if not os.path.isfile(path / 'models' / 'spm.model') or not os.path.isfile(path / 'models' / f'itos_{name}.pkl'):
-        raw_text = reduce(lambda t, rule: rule(t), rules, '\n'.join(text))
+        raw_text = reduce(lambda t, rule: rule(t), pre_rules, '\n'.join(text))
         raw_text_path = path / cache_name / 'all_text.txt'
         with open(raw_text_path, 'w') as f:
             f.write(raw_text)
@@ -86,18 +86,18 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N
                     f"--model_prefix={path / 'models' / 'spm'} " \
                     f"--vocab_size={vocab_size} --model_type={model_type} "
         spm.SentencePieceTrainer.Train(sp_params)
-  
+
         with open(path / 'models' / 'spm.vocab', 'r') as f:
             vocab = [line.split('\t')[0] for line in f.readlines()]
             vocab[0] = UNK
             vocab[pad_idx] = PAD
   
         pickle.dump(vocab, open(path / 'models' / f'itos_{name}.pkl', 'wb'))
-    
+    # todo add post rules
     vocab = Vocab(pickle.load(open(path / 'models' / f'itos_{name}.pkl', 'rb')))
     # We cannot use lambdas or local methods here, since `tok_func` needs to be
     # pickle-able in order to be called in subprocesses when multithread tokenizing
-    tokenizer = Tokenizer(tok_func=SentencepieceTokenizer, lang=str(path / 'models'), rules=rules)
+    tokenizer = Tokenizer(tok_func=SentencepieceTokenizer, lang=str(path / 'models'), pre_rules=pre_rules, post_rules=post_rules)
     
     clear_cache_directory(path, cache_name)
 
@@ -127,7 +127,7 @@ def ensure_paths_exists(*paths, message="One or more required files cannot be fo
     if error:
         raise FileNotFoundError(message)
 
-def get_data_folder():
+def get_data_folder() -> Path:
     """
     return data folder to use for future processing
     """
@@ -317,7 +317,8 @@ def read_clas_data(dir_path, dataset, lang) -> Tuple[Dict[str, List[List[str]]],
         # 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)
+        val_len = max(int(len(toks[TRN]) * 0.1), 2) # fastai does not work with validation set of size 1
+        trn_len = len(toks[TRN]) - val_len
         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:

tests/test_end_to_end.py

@@ -1,52 +1,55 @@
 import os
 import glob
 import fire
-import pytest
-
 import ulmfit.pretrain_lm
 import ulmfit.train_clas
 from fastai import *
 from fastai.text import *
-from fastai_contrib.utils 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 delete_test_models():
-    wt2 = data / 'wiki' / 'wikitext-2'
-    imdb = data / 'imdb'
-
-    # delete test models from the pretraining step
-    for test_file in glob.iglob(f'{str(wt2)}/models/end-to-end-test*'):
-        if os.path.isfile(test_file): os.remove(test_file)
-
-    # delete test vocab and model of sentencepiece training
-    for test_file in [wt2 / 'models' / 'spm.model', 
-                      wt2 / 'models' / 'spm.vocab']:
-        if os.path.isfile(test_file): os.remove(test_file)
-
-    # delete test models from the finetuning/classifier training step
-    for test_file in glob.iglob(f'{str(imdb)}/models/end-to-end-test*'):
-        if os.path.isfile(test_file): os.remove(test_file)
-
-
-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 copy_head(src_fn, dst_fn, n=1000):
+    with src_fn.open("r") as s, dst_fn.open("w") as d:
+        for i in range(n):
+            d.write(s.readline())
+
+
+def get_test_data():
+    data = get_data_folder()
+    wt = data / "wiki" / "wikitext-2"
+    imdb = data / "imdb"
+
+    test_data = data / "test"
+    shutil.rmtree(test_data)
+
+    test_wt = test_data / 'wikitext-s'
+    test_imdb = test_data / 'imdb'
+    test_wt.mkdir(exist_ok=True, parents=True)
+    test_imdb.mkdir(exist_ok=True, parents=True)
+
+    sz=10
+    # we use the same text to see if models can overfit
+    copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.train.tokens', n=10*sz)
+    copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.valid.tokens', n=6*sz)
+    copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.test.tokens', n=6*sz)
+    copy_head(imdb / 'train.csv', test_imdb / 'train.csv', n=10*sz)
+    copy_head(imdb / 'train.csv', test_imdb / 'test.csv', n=6*sz)
+
+    return test_data, test_wt
 
 
 def test_ulmfit_default_end_to_end():
     """  Test ulmfit with (default) Moses tokenizer on small wikipedia dataset.
     """
-    imdb, wt2 = check_data_exists()
+    test_data, wt2 = get_test_data()
     lm_name = 'end-to-end-test-default'
     cuda_id = 0
     results = ulmfit.pretrain_lm.pretrain_lm(
@@ -56,23 +59,20 @@ def test_ulmfit_default_end_to_end():
         qrnn=True,
         subword=False,
         max_vocab=1000,
-        bs=80,
+        bs=2,
         num_epochs=1,
-        name=lm_name,
-        ds_pct=0.03
-    )
-    assert results['accuracy'] > 0.30
+        name=lm_name)
+    assert results['accuracy'] > 0.02
 
     results = ulmfit.train_clas.new_train_clas(
-                data_dir=get_data_folder(),
-                lang='en', pretrain_name=lm_name, model_dir=wt2/'models',
-                qrnn=True,
-                cuda_id=cuda_id,
-                fine_tune=True,
-                max_vocab=1000,
-                bs=20, bptt=70, name=lm_name+'-imdb-clas',
-                dataset='imdb',
-                ds_pct=0.03)
+        data_dir=test_data,
+        lang='en', pretrain_name=lm_name, model_dir=wt2 / 'models',
+        qrnn=True,
+        cuda_id=cuda_id,
+        fine_tune=True,
+        max_vocab=1000,
+        bs=2, bptt=70, name=lm_name + '-imdb-clas',
+        dataset='imdb')
 
     delete_test_models()
 
@@ -80,7 +80,7 @@ def test_ulmfit_default_end_to_end():
 def test_ulmfit_sentencepiece_end_to_end():
     """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
     """
-    imdb, wt2 = check_data_exists()
+    imdb, wt2 = get_test_data()
     lm_name = 'end-to-end-test-spm'
     cuda_id = 0
     results = ulmfit.pretrain_lm.pretrain_lm(
@@ -101,3 +101,7 @@ def test_ulmfit_sentencepiece_end_to_end():
     #       sentencepiece for finetuning/classification is currently not implemented
 
     delete_test_models()
+
+
+if __name__ == "__main__":
+    fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz

ulmfit/train_clas.py

@@ -51,7 +51,7 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
         'Error: IMDb is only available in English.'
 
     data_dir = Path(data_dir)
-    assert data_dir.name == 'data',\
+    assert data_dir.name in ['data', 'test'],\
         f'Error: Name of data directory should be data, not {data_dir.name}.'
     dataset_dir = data_dir / dataset
     model_dir = Path(model_dir)
@@ -112,8 +112,9 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
                                   qrnn=qrnn, emb_sz=emb_sz, nh=nh, nl=nl, drop_mult=0.5)
 
     try:
+        print(f"Loading classifier {model_name}_{name}")
         learn.load(f'{model_name}_{name}')
-        print("Loading classifier")
+
     except FileNotFoundError:
         learn.load_encoder(lm_enc_finetuned)
         print("loading encoder")
@@ -147,9 +148,8 @@ def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_typ
     if not (tmp_dir / f'{TRN}_{lang}_ids.npy').exists():
         print('Reading the data...')
         toks, lbls = read_clas_data(dataset_dir, dataset, lang)
-
         # create the vocabulary
-        counter = Counter(word for example in np.concatenate([toks[TRN],toks[TST],toks[VAL]]) for word in example)
+        counter = Counter(word for example in toks[TRN]+toks[TST]+toks[VAL] for word in example)
         itos = [word for word, count in counter.most_common(n=max_vocab)]
         itos.insert(0, PAD)
         itos.insert(0, UNK)
@@ -176,14 +176,17 @@ def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_typ
           f'Test size: {len(ids[TST])}.')
     if ds_pct < 1.0:
         print(f"Making the dataset smaller {ds_pct}")
-        for split in [TRN, VAL, TST]:
-            ids[split] = ids[split][:int(len(ids[split]) * ds_pct)]
+    for split in [TRN, VAL, TST]:
+        ids[split] = np.array([np.array(e, dtype=np.int) for e in ids[split]])
+        lbls[split] = np.array([np.array(e, dtype=np.int) for e in lbls[split]])
     data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=np.concatenate([ids[TRN],ids[TST]]),
                                        valid_ids=ids[VAL], bs=bs, bptt=bptt, lm_type=lm_type)
     #  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, classes={l:l for l in lbls[VAL]})
+        train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs, classes={l:l for l in lbls[TRN]})
+
+    print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
     return data_clas, data_lm