Merge pull request #18 from n-waves/refactor

[WIP] Refactoring #17

README.md

@@ -1,6 +1,18 @@
 # ulmfit-multilingual
 Temporary repository used for collaboration on application of for multiple languages.
 
+# How to train classifier
+
+```
+$ python -m ulmfit lm --dataset-path data/wiki/wikitext-103 --bidir=False --qrnn=False --tokenizer=vf --name 'bs40' --bs=40 --cuda-id=0  -  train 20 --drop-mult=0.9
+...
+Model dir: data/wiki/wikitext-103/models/vf60k/lstm_bs40.m
+...
+$ python -m ulmfit cls --dataset-path data/imdb --base-lm-path data/wiki/wikitext-103/models/vf60k/lstm_bs40.m - train 20   
+```
+ 
+
+
 ## data directory strucutre
 
 Directory structure after changes to the way we process wiki dumps.

fastai_contrib/data.py

@@ -13,23 +13,24 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
     "Create a dataloader with bptt slightly changing."
     def __init__(self, dataset:LabelList, bs:int=64, bptt:int=70,
                  lm_type:LanguageModelType=LanguageModelType.FwdLM, shuffle:bool=False,
-                 max_len:int=25):
-        self.dataset,self.bs,self.bptt,self.lm_type,self.shuffle = dataset,bs,bptt,lm_type,shuffle
+                 max_len:int=25, p_bptt:int=0.95):
+        self.dataset,self.bs,self.bptt,self.lm_type,self.shuffle, self.p_bptt = dataset,bs,bptt,lm_type,shuffle,p_bptt
         self.first,self.i,self.iter = True,0,0
-        self.n = len(np.concatenate(dataset.x.items)) // self.bs
+        self.n = len(np.concatenate(dataset.x.items)) // self.bs if len(dataset.x.items) > 0 else 0
         self.max_len,self.num_workers = max_len,0
+        self.init_kwargs = dict(bs=bs, bptt=bptt, lm_type=lm_type, shuffle=shuffle, max_len=max_len, p_bptt=p_bptt)
 
     def __iter__(self):
         if getattr(self.dataset, 'item', None) is not None:
-            yield LongTensor(getattr(self.dataset, 'item')).unsqueeze(1),LongTensor([0])
+            yield LongTensor(getattr(self.dataset, 'item'))[None],LongTensor([0])
         idx = np.random.permutation(len(self.dataset)) if self.shuffle else range(len(self.dataset))
-        data = self.batchify(np.concatenate([np.array(self.dataset.x.items[i], dtype=np.int) for i in idx]))
+        data = self.batchify(np.concatenate([self.dataset.x.items[i] for i in idx]))
 
         pos, itr = 0,0
         while pos < self.n-1 and itr<len(self):
             if self.first and pos == 0: self.first,seq_len = False,self.bptt + self.max_len
             else:
-                bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2.
+                bptt = self.bptt if np.random.random() < self.p_bptt else self.bptt / 2.
                 seq_len = max(5, int(np.random.normal(bptt, 5)))
                 seq_len = min(seq_len, self.bptt + self.max_len)
             res = self.get_batch(data, pos, seq_len)
@@ -41,27 +42,24 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
     def __getattr__(self,k:str)->Any: return getattr(self.dataset, k)
 
     @property
-    def batch_size(self):
-        return self.bs
-
+    def batch_size(self): return self.bs
     @batch_size.setter
-    def batch_size(self, v):
-        self.bs = v
+    def batch_size(self, v): self.bs = v
 
     def batchify(self, data:np.ndarray) -> LongTensor:
         "Split the corpus `data` in batches."
         nb = data.shape[0] // self.bs
-        data = np.array(data[:nb*self.bs]).reshape(self.bs, -1).T
-        if self.lm_type == LanguageModelType.BwdLM: data=data[::-1].copy()
-        elif self.lm_type == LanguageModelType.BiLM: data = np.stack([data, data[::-1].copy()], axis=2)
+        data = np.array(data[:nb*self.bs]).reshape(self.bs, -1)
+        if self.lm_type == LanguageModelType.BwdLM: data = data[:,::-1].copy()
+        elif self.lm_type == LanguageModelType.BiLM: data = np.stack([data, data[:,::-1].copy()], axis=2)
         return LongTensor(data)
 
     def get_batch(self, data:LongTensor,  i:int, seq_len:int) -> Tuple[LongTensor, LongTensor]:
         "Create a batch at `i` of a given `seq_len`."
-        seq_len = min(seq_len, len(data) - 1 - i)
-        x = data[i:i+seq_len]
-        y = data[i+1:i+1+seq_len].contiguous() # x & y has 2 elements on the last dimension
-        y = y.view(-1, 2) if self.lm_type == LanguageModelType.BiLM else y.view(-1)
+        seq_len = min(seq_len, data.shape[1] - 1 - i)
+        x = data[:,i:i+seq_len]
+        y = data[:,i+1:i+1+seq_len]
+        y = y.contiguous().view(-1, 2) if self.lm_type == LanguageModelType.BiLM else y.contiguous().view(-1)
         return x,y
 
 #endregion

fastai_contrib/learner.py

@@ -1,11 +1,11 @@
-from fastai import GradientClipping, accuracy
-from fastai.callbacks import *
-from fastai.basic_data import *
-from fastai.datasets import untar_data
-from fastai_contrib.models import get_bilm, get_rnn_classifier, get_birnn_classifier
-from fastai.text.learner import *
+from torch.nn import CrossEntropyLoss
+
+from fastai import *
+from fastai.text import *
 
 #region New code
+from fastai_contrib.models import *
+
 
 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,
@@ -25,23 +25,28 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in
         fnames = [learn.path/learn.model_dir/f'{fn}.{ext}' for fn,ext in zip(pretrained_fnames, ['pth', 'pkl'])]
         learn.load_pretrained(*fnames)
         learn.freeze()
+    learn.loss_func = CrossEntropyLoss()  # I'm not sure why fast ai is using CrossEntropyFlat but it breaks bilm
     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':
+                            drop_mult: float = 1., qrnn: bool = False, bicls_head:str='BiPoolingLinearClassifier',  **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]
+    if bicls_head == 'BiPoolingLinearClassifier':
+        count = 3*2
+    else:
+        count = 3
+    layers = [emb_sz * count] + 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)
+                               qrnn=qrnn, bicls_head=bicls_head)
     learn = RNNLearner(data, model, bptt, split_func=birnn_classifier_split, **kwargs)
     return learn
 
@@ -78,18 +83,25 @@ def convert_weights(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:Collection[s
 
 def convert_weights_with_prefix(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:Collection[str], prefix='') -> Weights:
     "Convert the model weights to go with a new vocabulary."
-    dec_bias, enc_wgts = wgts[prefix+'1.decoder.bias'], wgts[prefix+'0.encoder.weight']
-    bias_m, wgts_m = dec_bias.mean(0), enc_wgts.mean(0)
-    new_w = enc_wgts.new_zeros((len(itos_new),enc_wgts.size(1))).zero_()
-    new_b = dec_bias.new_zeros((len(itos_new),)).zero_()
-    for i,w in enumerate(itos_new):
-        r = stoi_wgts[w] if w in stoi_wgts else -1
-        new_w[i] = enc_wgts[r] if r>=0 else wgts_m
-        new_b[i] = dec_bias[r] if r>=0 else bias_m
-    wgts[prefix+'0.encoder.weight'] = new_w
-    wgts[prefix+'0.encoder_dp.emb.weight'] = new_w.clone()
-    wgts[prefix+'1.decoder.weight'] = new_w.clone()
-    wgts[prefix+'1.decoder.bias'] = new_b
+    if 'model' in wgts:
+        wgts['model'] = convert_weights_with_prefix(wgts['model'], stoi_wgts, itos_new, prefix)
+    else:
+        dec_bias, enc_wgts = wgts[prefix+'1.decoder.bias'], wgts[prefix+'0.encoder.weight']
+        bias_m, wgts_m = dec_bias.mean(0), enc_wgts.mean(0)
+        new_w = enc_wgts.new_zeros((len(itos_new),enc_wgts.size(1))).zero_()
+        new_b = dec_bias.new_zeros((len(itos_new),)).zero_()
+        unk_tokens=[]
+        for i,w in enumerate(itos_new):
+            r = stoi_wgts[w] if w in stoi_wgts else -1
+            if r < 0:
+                unk_tokens.append(w)
+            new_w[i] = enc_wgts[r] if r>=0 else wgts_m
+            new_b[i] = dec_bias[r] if r>=0 else bias_m
+        print(f"Unknown tokens {len(unk_tokens)}, first 100: {unk_tokens[:100]}")
+        wgts[prefix+'0.encoder.weight'] = new_w
+        wgts[prefix+'0.encoder_dp.emb.weight'] = new_w.clone()
+        wgts[prefix+'1.decoder.weight'] = new_w.clone()
+        wgts[prefix+'1.decoder.bias'] = new_b
     return wgts
 
 #endregion

fastai_contrib/models.py

@@ -6,10 +6,11 @@ from fastai.text.models import *
 
 class BiLMModel(nn.Module):
 
-    def __init__(self, fwd_lm:nn.Module, bwd_lm:nn.Module):
+    def __init__(self, fwd_lm:nn.Module, bwd_lm:nn.Module, squash_bs_sl=False):
         super().__init__()
         self.fwd_lm = fwd_lm
         self.bwd_lm = bwd_lm
+        self.squash_bs_sl = squash_bs_sl
 
     def __getitem__(self, idx):
         return BiLMModel(self.fwd_lm[idx], self.bwd_lm[idx])
@@ -29,55 +30,63 @@ class BiLMModel(nn.Module):
             b = input[..., 1]
         elif len(input.shape) == 2: # sl, bs - support during classification mode
             f = input
-            b = torch.flip(input, [0])
+            b = torch.flip(input, [1]) # todo test if we are duplicating the backward pass correctly
         else:
             raise AttributeError(f"Inorrect size of input, {input.shape}")
         fwd_o = self.fwd_lm(f)
         bwd_o = self.bwd_lm(b)
 
-        return self.stack(fwd_o, bwd_o)
+        outs = self.stack(fwd_o, bwd_o)
+        if self.squash_bs_sl:
+            o = outs[0]
+            o = o.view(o.shape[0]*o.shape[1],o.shape[2],o.shape[3])
+            outs[0] = o
+        return outs
 
     def reset(self):
         "Reset the hidden states of underlaying lms."
         self.fwd_lm.reset()
         self.bwd_lm.reset()
 
+class MultiBatchBiLMModel(BiLMModel):
+    "Create a RNNCore module that can process a full sentence."
 
-class BiPoolingLinearClassifier(nn.Module):
+    def __init__(self, bptt:int, max_seq:int, *args, **kwargs):
+        self.max_seq,self.bptt = max_seq,bptt
+        super().__init__(*args, **kwargs)
+
+    def concat(self, arrs:Collection[Tensor])->Tensor:
+        "Concatenate the `arrs` along the batch dimension."
+        return [torch.cat([l[si] for l in arrs], dim=1) for si in range_of(arrs[0])]
+
+    def forward(self, input:LongTensor)->Tuple[Tensor,Tensor]:
+        bs,sl = input.size()
+        self.reset()
+        raw_outputs, outputs = [],[]
+        for i in range(0, sl, self.bptt):
+            r, o = super().forward(input[:,i: min(i+self.bptt, sl)])
+            if i>(sl-self.max_seq):
+                raw_outputs.append(r)
+                outputs.append(o)
+        return self.concat(raw_outputs), self.concat(outputs)
+
+class BiPoolingLinearClassifier(PoolingLinearClassifier):
     "Create a linear classifier with pooling."
 
-    def __init__(self, layers:Collection[int], drops:Collection[float]):
-        super().__init__()
-        mod_layers = []
-        activs = [nn.ReLU(inplace=True)] * (len(layers) - 2) + [None]
-        for n_in,n_out,p,actn in zip(layers[:-1],layers[1:], drops, activs):
-            mod_layers += bn_drop_lin(n_in, n_out, p=p, actn=actn)
-        self.layers = nn.Sequential(*mod_layers)
-
-    def pool(self, x:Tensor, bs:int, is_max:bool):
-        "Pool the tensor along the seq_len dimension."
-        f = F.adaptive_max_pool1d if is_max else F.adaptive_avg_pool1d
-        return f(x.permute(1,2,0), (1,)).view(bs,-1)
-
     def forward(self, input:Tuple[Tensor,Tensor])->Tuple[Tensor,Tensor,Tensor]:
         raw_outputs, outputs = input
         output = outputs[-1]
         if len(output.size()) == 3:
-            sl,bs,_ = output.size()
-            avgpool = self.pool(output, bs, False)
-            mxpool = self.pool(output, bs, True)
-            x = torch.cat([output[-1], mxpool, avgpool], 1)
-            x = self.layers(x)
-            return x, raw_outputs, outputs
+            return super().forward(input)
         elif len(output.size()) == 4:
-            sl, bs, em_sz, passes = output.size()
+            bs, sl, em_sz, passes = output.size()
 
             f_avgpool = self.pool(output[...,0], bs, False)
             f_mxpool = self.pool(output[...,0], bs, True)
             b_avgpool = self.pool(output[..., 1], bs, False)
             b_mxpool = self.pool(output[..., 1], bs, True)
-            x = torch.cat([output[-1][..., 0], f_mxpool, f_avgpool,
-                           output[-1][..., 1], b_mxpool, b_avgpool,], 1)
+            x = torch.cat([output[:,-1,..., 0], f_mxpool, f_avgpool,
+                           output[:,-1,..., 1], b_mxpool, b_avgpool,], 1)
             x = self.layers(x)
             return x, raw_outputs, outputs
 
@@ -134,18 +143,23 @@ def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, t
 
     return BiLMModel(
         fwd_lm=SequentialRNN(fwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)),
-        bwd_lm=SequentialRNN(bwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)))
+        bwd_lm=SequentialRNN(bwd_rnn_enc, LinearDecoder(vocab_sz, emb_sz, output_p, tie_encoder=enc, bias=bias)),
+        squash_bs_sl=True)
 
 def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int,
                        pad_token:int, layers:Collection[int], drops:Collection[float], bidir:bool=False, qrnn:bool=False,
-                       hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5)->nn.Module:
+                       hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5, bicls_head:str='BiPoolingLinearClassifier')->nn.Module:
     "Create a RNN classifier model."
     fwd_rnn_enc = MultiBatchRNNCore(bptt, max_seq, vocab_sz, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                       qrnn=qrnn, hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
     bwd_rnn_enc = MultiBatchRNNCore(bptt, max_seq, vocab_sz, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                                 qrnn=qrnn, hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
 
-    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), AvgPoolingLinearClassifier(layers, drops))
+    head = BiPoolingLinearClassifier
+    if bicls_head == 'BiPoolingLinearClassifier': head = BiPoolingLinearClassifier
+    elif bicls_head == 'AvgPoolingLinearClassifier': head = AvgPoolingLinearClassifier
+
+    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), head(layers, drops))
     model.reset()
     return model
 

fastai_contrib/utils.py

@@ -1,18 +1,9 @@
 """
 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.transform import Tokenizer, BaseTokenizer, Vocab
-from fastai.torch_core import *
+from fastai import *
+from fastai.text import *
 
 import shutil
 import pathlib
@@ -21,10 +12,15 @@ 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
+EOS = 'xxeos' # fastai does not use eos, but we do
+SEP = 'xxsep' # special separator token for NLI
+
+def replace_std_toks(x:str) -> str:
+    "Replace standard token names with fastai supported tokens"
+    # We change tokens to f'xx{token_name}' as it is not split by Moses tokenizer,
+    # while f'<{token_name}>' is being split to: '<' f'{token_name}' '>'
+    return x.replace('<unk>', UNK).replace('<bos>', BOS).replace('<eos>', EOS)
+
 PAD_TOKEN_ID = 1
 IMDB, XNLI, TRN, VAL, TST, EN = 'imdb', 'xnli', 'train', 'val', 'test', 'en'
 DATASETS = ['imdb', 'xnli']
@@ -39,68 +35,103 @@ CLASSES = ['neg', 'pos', 'unsup']
 number_match_re = re.compile(r'^([0-9]+[,.]?)+$')
 number_split_re = re.compile(r'([,.])')
 
-class SentencepieceTokenizer(BaseTokenizer):
-    def __init__(self, model_dir:PathOrStr):
+class MosesTokenizerFunc(BaseTokenizer):
+    "Wrapper around a MosesTokenizer to make it a `BaseTokenizer`."
+    def __init__(self, lang:str):
+        super().__init__(lang=lang)
+        self.tok = MosesTokenizer(lang)
+
+    def tokenizer(self, t:str) -> List[str]:
+        return self.tok.tokenize(t, return_str=False, escape=False)
+
+    def add_special_cases(self, toks:Collection[str]):
+        for w in toks:
+            assert len(self.tokenizer(w))==1, f"Tokenizer is unable to keep {w} as one token!"
+
+class SentencePieceTokenizer(Tokenizer):
+    "Put together rules and a tokenizer function to tokenize text with multiprocessing."
+    def __init__(self, spm_model, lang:str='en', pre_rules:ListRules=None,
+                 post_rules:ListRules=None, special_cases:Collection[str]=None, n_cpus:int=None, use_moses=False):
+        super().__init__(self.tok_fun_with_sp, lang, pre_rules, post_rules, special_cases, n_cpus)
+        self.spm_model = spm_model
+        self.use_moses = use_moses
+
+    def tok_fun_with_sp(self, lang):
         try:
-            import sentencepiece as spm  
+            import sentencepiece as spm
         except ImportError:
             raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')
-        self.tok = spm.SentencePieceProcessor()
-        self.tok.Load(str(pathlib.Path(model_dir) / 'spm.model'))
-    
-    def tokenizer(self, t:str) -> List[str]:
-        return self.tok.EncodeAsPieces(t)
-    
-    def add_special_cases(self, toks:Collection[str]):
-        pass
+        tok = MosesTokenizerFunc(lang) if self.use_moses else BaseTokenizer(lang)
+        tok.sp = spm.SentencePieceProcessor()
+        tok.sp.Load(str(self.spm_model))
+        return tok
 
+    def process_text(self, t:str, tok:BaseTokenizer) -> List[str]:
+        "Process one text `t` with tokenizer `tok`."
+        toks = super().process_text(t, tok)
+        toks = tok.sp.EncodeAsPieces(" ".join(toks))
+        return toks
 
-def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, pre_rules:ListRules=None, post_rules:ListRules=None,
+def get_sentencepiece(cache_dir:PathOrStr, load_text,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):
+                      use_moses=False, lang='en'):
     try:
         import sentencepiece as spm
     except ImportError:
         raise Exception('sentencepiece module is missing: run `pip install sentencepiece`')
-    
-    path = pathlib.Path(path)
-    cache_name = 'tmp'
-    os.makedirs(path / cache_name, exist_ok=True)
-    os.makedirs(path / 'models', exist_ok=True)
-    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), 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)
-      
-        sp_params = f"--input={raw_text_path} --pad_id={pad_idx} --unk_id=0 " \
-                    f"--character_coverage=1.0 --bos_id=-1 --eos_id=-1 " \
-                    f"--input_sentence_size={int(input_sentence_size)} " \
-                    f"--model_prefix={path / 'models' / 'spm'} " \
-                    f"--vocab_size={vocab_size} --model_type={model_type} "
-        spm.SentencePieceTrainer.Train(sp_params)
+    cache_dir = pathlib.Path(cache_dir)
+    pre_rules = pre_rules if pre_rules is not None else defaults.text_pre_rules
+    post_rules = post_rules if post_rules is not None else defaults.text_post_rules
 
-        with open(path / 'models' / 'spm.vocab', 'r') as f:
+    special_cases = defaults.text_spec_tok
+
+    if not os.path.isfile(cache_dir / 'spm.model') or not os.path.isfile(cache_dir / f'itos.pkl'):
+        # load the text from the train tokens file
+        text = load_text()
+        text = filter(lambda x: len(x.rstrip(" ")), text)
+        text = (reduce(lambda t, rule: rule(t), pre_rules, line) for line in text)
+        if use_moses:
+            mt = MosesTokenizer(lang)
+            splitter = lambda t: mt.tokenize(t, return_str=False, escape=False)
+        else:
+            splitter = lambda t: t.split()
+        def cleanup_n_postprocess(t):
+            t = splitter(t)
+            for r in post_rules:
+                t = r(t)
+            return ' '.join(t)
+        text = map(cleanup_n_postprocess, text)
+        raw_text_path = cache_dir / 'all_text.txt'
+        with open(raw_text_path, 'w') as f: f.write("\n".join(text))
+
+        sp_params = [
+            f"--input={raw_text_path}",
+            f"--character_coverage=1.0",
+            f"--unk_id={len(defaults.text_spec_tok)}",
+            f"--pad_id=-1",
+            f"--bos_id=-1",
+            f"--eos_id=-1",
+            f"--max_sentence_length=20480",
+            f"--input_sentence_size={int(input_sentence_size)}",
+            f"--user_defined_symbols={','.join(special_cases)}",
+            f"--model_prefix={cache_dir/'spm'}",
+            f"--vocab_size={vocab_size} --model_type={model_type}"]
+        spm.SentencePieceTrainer.Train(" ".join(sp_params))
+
+        with open(cache_dir / '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'))
+
+        pickle.dump(vocab, open(cache_dir/ f'itos.pkl', 'wb'))
     # todo add post rules
-    vocab = Vocab(pickle.load(open(path / 'models' / f'itos_{name}.pkl', 'rb')))
+    vocab = Vocab(pickle.load(open(cache_dir / f'itos.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'), pre_rules=pre_rules, post_rules=post_rules)
-    
-    clear_cache_directory(path, cache_name)
-
+    tokenizer = SentencePieceTokenizer(cache_dir/'spm.model',
+                                       use_moses=use_moses,
+                                       lang=lang,
+                                       pre_rules=pre_rules,
+                                       post_rules=post_rules)
     return {'tokenizer': tokenizer, 'vocab': vocab}
 
 
@@ -108,7 +139,6 @@ def clear_cache_directory(path:PathOrStr, cache_name:str='tmp'):
     path = pathlib.Path(path)
     shutil.rmtree(path / cache_name)
 
-
 def get_texts(path):
     texts, labels = [],[]
     for idx, label in enumerate(CLASSES):
@@ -187,63 +217,9 @@ def prepare_imdb(file_path: str, prepare_lm = False):
     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)
-
+    df_trn[df_trn['labels'] == 2].to_csv(CLAS_PATH / 'unsup.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, spm_path=None) -> 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)
-    :param spm_path: path to sentencepiece model
-    :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')
-    if spm_path is not None:
-        sp = SentencepieceTokenizer(spm_path)
-
-    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(int(label))
-            raw_tokens = mt.tokenize(text, return_str=True).split(' ')
-            
-            tokens = []
-            
-            # fix up occurences of numbers in text
-            for token in raw_tokens:
-                if number_match_re.match(token):
-                    tokens += number_split_re.sub(r' @\1@ ', token).split()
-                else:
-                    tokens.append(token)
-            
-            if spm_path is not None:
-                tokens = sp.tokenizer(' '.join(tokens))
-
-            toks.append(tokens + [EOS])
-    return toks, lbls
-
-
 def read_xnli(dir_path, lang, split, spm_path=None) -> Tuple[List[List[str]], List[str]]:
     """
     Reads XNLI data.
@@ -262,7 +238,14 @@ def read_xnli(dir_path, lang, split, spm_path=None) -> Tuple[List[List[str]], Li
     file_path = dir_path / file_path
     
     if spm_path is not None:
-        sp = SentencepieceTokenizer(spm_path)
+        tokenizer = SentencePieceTokenizer(spm_path,
+                                           use_moses=False,
+                                           lang=lang)
+        tok = tokenizer.tok_fun_with_sp(lang)
+        tokenize = lambda x: tokenizer.process_text(x, tok)
+        print("WARNING: Sentence Piece is not tested on XNLI yet")
+    else:
+        tokenize = lambda x: x.split(' ')
 
     toks, lbls = [], []
     print(f'Reading {file_path}...')
@@ -281,13 +264,9 @@ def read_xnli(dir_path, lang, split, spm_path=None) -> Tuple[List[List[str]], Li
                 premise, hypo, label = row[-3], row[-2], row[1]
             
             # TODO add BOS
-            if spm_path is not None:
-                premise_toks = sp.tokenizer(premise) + [EOS]
-                hypo_toks = sp.tokenizer(hypo) + [EOS]
-            else:
-                premise_toks = premise.split(' ') + [EOS]
-                hypo_toks = hypo.split(' ') + [EOS]
-            
+            premise_toks = tokenize(premise) + [EOS]
+            hypo_toks = tokenize(hypo) + [EOS]
+
             toks.append(premise_toks + [SEP] + hypo_toks)
             lbls.append(label)
     return toks, lbls
@@ -304,7 +283,6 @@ def read_clas_data(dir_path, dataset, lang) -> Tuple[Dict[str, List[List[str]]],
              2. a dictionary mapping splits to a list of labels
     """
     processors = {
-        'imdb': read_imdb,
         'xnli': read_xnli
     }
     processor = processors[dataset]
@@ -332,14 +310,17 @@ def replace_number(token):
     return token
 
 
-def read_file(file_path, outname):
+def read_file(file_path, outname=None):
     """Reads a text file and writes it to a .csv."""
     with open(file_path, encoding='utf8') as f:
         text = f.readlines()
     df = pd.DataFrame(
         {'text': np.array(text), 'labels': np.zeros(len(text))},
         columns=['labels', 'text'])
-    df.to_csv(file_path.parent / f'{outname}.csv', header=False, index=False)
+    if outname is not None:
+        df.to_csv(file_path.parent / f'{outname}.csv', header=False, index=False)
+    return df
+
 
 
 def read_whitespace_file(filepath):
@@ -351,9 +332,6 @@ def read_whitespace_file(filepath):
             tokens.append(line.split() + [EOS])
     return np.array(tokens)
 
-
-
-
 class DataStump:
     """Placeholder class as LanguageModelLoader requires object with ids attribute."""
     def __init__(self, ids):

prepare_imdb.sh

@@ -6,6 +6,6 @@ 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
+echo "Imdb is raw text no preparation is done"
+python -m fastai_contrib.utils prepare_imdb "${DATA_DIR}/aclImdb.tgz"
 

results/sp30k_biqrnn_bs70

@@ -0,0 +1,13 @@
+Starting from random weights                                                                                                                                             
+epoch  train_loss  valid_loss  accuracy_fwd  accuracy_bwd                                                                                                                
+1      3.669261    3.641705    0.399714      0.380689                                                                                                                    
+2      3.574335    3.547273    0.404729      0.385104                                                                                                                    
+3      3.573150    3.549644    0.403350      0.384167                                                                                                                    
+4      3.518714    3.499090    0.408166      0.389015
+5      3.477355    3.441828    0.413880      0.394777
+6      3.408005    3.366269    0.422041      0.402934
+7      3.314280    3.284519    0.431068      0.411727
+8      3.244735    3.205757    0.440180      0.421078
+9      3.170936    3.152495    0.446947      0.428045
+10     3.131996    3.138446    0.448782      0.430013
+Saving optimiser state at data/wiki/wikitext-103/models/sp30k/biqrnn_bs70.m

tests/test_end_to_end.py

@@ -12,14 +12,13 @@ It is a mixture of a pytest unit test and woven together to compose an end to en
 """
 
 import fastai.core
-fastai.core.turn_off_parallel_execution=True
-
+fastai.core.defaults.cpus = 1
+cuda_id=0
 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"
@@ -35,71 +34,151 @@ def get_test_data():
 
     sz=1
     # 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(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.train.tokens', n=1000*sz)
+    copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.valid.tokens', n=600*sz)
+    copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.test.tokens', n=600*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)
+    copy_head(imdb / 'train.csv', test_imdb / 'unsup.csv', n=1*sz)
 
     return test_data, test_wt
 
 
+def test_ulmfit_works_with_relative_paths():
+    """  Test ulmfit with (default) Moses tokenizer on small wikipedia dataset.
+    """
+    os.chdir(get_data_folder()/"..")
+
+
+    test_data, wt2 = get_test_data()
+    lm_name = 'end-to-end-test-default'
+    cuda_id = 0
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2.relative_to(Path.cwd()),
+        lang='en',
+        qrnn=False,
+        max_vocab=1000,
+        name=lm_name,
+        cuda_id=cuda_id)
+
+    exp.train_lm(num_epochs=1, bs=2)
+
+    #assert exp.results['accuracy'] > 0.02
+
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=1, unfreeze=False, bs=4,)
+
+    # should work for the second time as well
+
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, )
+
+
 def test_ulmfit_default_end_to_end():
     """  Test ulmfit with (default) Moses tokenizer on small wikipedia dataset.
     """
     test_data, wt2 = get_test_data()
     lm_name = 'end-to-end-test-default'
     cuda_id = 0
-    results = ulmfit.pretrain_lm.pretrain_lm(
-        dir_path=wt2,
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2,
+        lang='en',
+        qrnn=False,
+        max_vocab=1000,
+        name=lm_name,
+        cuda_id=cuda_id)
+
+    exp.train_lm(num_epochs=1, bs=2)
+
+    #assert exp.results['accuracy'] > 0.02
+
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4,)
+
+def test_ulmfit_fastai_end_to_end():
+    """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
+    """
+    test_data, wt2 = get_test_data()
+    lm_name = 'end-to-end-test-fastai'
+
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
-        subword=False,
-        max_vocab=1000,
-        bs=2,
-        num_epochs=1,
-        name=lm_name)
-    assert results['accuracy'] > 0.02
+        qrnn=False,
+        tokenizer='f',
+        max_vocab=100,
+        name=lm_name,
+    )
+    exp.train_lm(num_epochs=1, bs=2)
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, )
 
-    results = ulmfit.train_clas.new_train_clas(
-        data_dir=test_data,
-        lang='en', pretrain_name=lm_name, model_dir=wt2 / 'models',
-        qrnn=True,
+def test_ulmfit_fastai_bidir_end_to_end():
+    """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
+    """
+    test_data, wt2 = get_test_data()
+    lm_name = 'end-to-end-test-fastai'
+
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2,
+        lang='en',
         cuda_id=cuda_id,
-        fine_tune=True,
-        max_vocab=1000,
-        num_lm_epochs=0,
-        bs=4, # minimum size is 4 otherwise it somewhere becomes 1 and fit stops working
-        bptt=70,
-        name=lm_name + '-imdb-clas',
-        dataset='imdb')
+        qrnn=False,
+        bidir=True,
+        tokenizer='f',
+        max_vocab=100,
+        name=lm_name,
+    )
+    exp.train_lm(num_epochs=1, bs=2)
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(str(test_data / 'imdb'), str(exp.model_dir))
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, )
 
+def test_ulmfit_moses_fa_bidir_end_to_end():
+    """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
+    """
+    test_data, wt2 = get_test_data()
+    lm_name = 'end-to-end-test-fastai'
+
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2,
+        lang='en',
+        cuda_id=cuda_id,
+        qrnn=False,
+        bidir=True,
+        tokenizer='vf',
+        max_vocab=100,
+        name=lm_name,
+    )
+    exp.train_lm(num_epochs=1, bs=2)
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, )
+
+# def test_classification_model_work_with_different_dropmul():
+#     learn = self.create_cls_learner(data_clas, drop_mult=0.1)
+#     learn = self.create_cls_learner(data_clas, drop_mult=0.0)
 
 def test_ulmfit_sentencepiece_end_to_end():
     """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
     """
-    imdb, wt2 = get_test_data()
+    test_data, wt2 = get_test_data()
     lm_name = 'end-to-end-test-spm'
-    cuda_id = 0
-    results = ulmfit.pretrain_lm.pretrain_lm(
-        dir_path=wt2,
+
+    exp = ulmfit.pretrain_lm.LMHyperParams(
+        dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
-        subword=True,
-        max_vocab=100,
-        bs=2,
-        num_epochs=1,
+        qrnn=False,
+        tokenizer=ulmfit.pretrain_lm.Tokenizers.SUBWORD,
+        max_vocab=200,
         name=lm_name,
     )
-
-    assert results['accuracy'] > 0.30
-
-    # NOTE: ds_pct is not available for sentencepiece -- tests are on the complete dataset
-    #       sentencepiece for finetuning/classification is currently not implemented
+    exp.train_lm(num_epochs=1, bs=2)
+    # not supported yet
+    exp2 = ulmfit.train_clas.CLSHyperParams.from_lm(test_data / 'imdb', exp.model_dir)
+    exp2.train_cls(num_lm_epochs=0, unfreeze=False, bs=4, )
 
 
 if __name__ == "__main__":
-    fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz
+    fire.Fire()  # allows using all functions via CLI
 

ulmfit/__init__.py

@@ -0,0 +1 @@
+

ulmfit/__main__.py

@@ -0,0 +1,26 @@
+from functools import wraps
+
+import fire
+from .pretrain_lm import LMHyperParams
+from .train_clas import CLSHyperParams
+
+class FireView:
+    def __init__(self, **kwargs):
+        for k,v in kwargs.items():
+            setattr(self, k, v)
+
+class ULMFiT:
+    @wraps(LMHyperParams)
+    def lm(self, dataset_path,  **changes):
+        changes['dataset_path'] = dataset_path
+        params = LMHyperParams(**changes)
+        return FireView(train=params.train_lm)
+
+    lm2 = LMHyperParams
+    @wraps(CLSHyperParams)
+    def cls(self, dataset_path, base_lm_path, **changes):
+        params = CLSHyperParams.from_lm(dataset_path, base_lm_path, **changes)
+        return FireView(train=params.train_cls)
+
+if __name__ == '__main__':
+    fire.Fire(ULMFiT())

ulmfit/create_wikitext.py

@@ -24,7 +24,7 @@ def get_texts(root):
                     if text.strip() == title:
                         # print('No content continuing...')
                         continue
-                    yield text
+                    yield (f"={title}=\n"+text)
 
 
 def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'):

ulmfit/pretrain_lm.py

@@ -4,15 +4,19 @@ expected to have been tokenized with Moses and processed with `postprocess_wikit
 That is, the data is expected to be white-space separated and numbers are expected
 to be split.
 """
+from dataclasses import InitVar
+
 import fastai
 import fire
 
 from fastai import *
+from fastai.callbacks import CSVLogger, SaveModelCallback
 from fastai.text import *
 import torch
 from fastai_contrib.utils import read_file, read_whitespace_file, \
-    validate, PAD, UNK, get_sentencepiece
-from fastai_contrib.learner import bilm_learner, accuracy_fwd, accuracy_bwd
+    validate, PAD, UNK, get_sentencepiece, read_clas_data, TRN, VAL, TST, PAD_TOKEN_ID, MosesTokenizerFunc, \
+    replace_std_toks
+from fastai_contrib.learner import bilm_learner, accuracy_fwd, accuracy_bwd, bilm_text_classifier_learner
 import pickle
 
 from pathlib import Path
@@ -20,164 +24,237 @@ from pathlib import Path
 from collections import Counter
 import fastai_contrib.data as contrib_data
 
-# to install, do:
-# conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92]
-# cupy needs to be installed for QRNN
+LM_BEST = "lm_best"
+ENC_BEST = "enc_best"
 
 
-def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, 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 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 qrnn: Use a QRNN. Requires installing cupy.
-    :param subword: Use sub-word tokenization on the cleaned data.
-    :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 model_dir: The path to the directory where the models should be saved
-    :param bidir: whether the language model is bidirectional
-    """
-    results = {}
+class Tokenizers(Enum):
+    SUBWORD='sp'
+    MOSES='v'
+    MOSES_FA='vf'
+    FASTAI='f'
 
-    if not torch.cuda.is_available():
-        print('CUDA not available. Setting device=-1.')
-        cuda_id = -1
-    torch.cuda.set_device(cuda_id)
+def istitle(line):
+    return len(re.findall(r'^ ?= [^=]* = ?$', line)) != 0
 
-    dir_path = Path(dir_path)
-    assert dir_path.exists()
-    model_dir = dir_path / 'models'  # removed from params, as it is absolute models location in train_clas and here it is relative
-    model_dir.mkdir(exist_ok=True)
-    print('Batch size:', bs)
-    print('Max vocab:', max_vocab)
-    model_name = 'qrnn' if qrnn else 'lstm'
-    if qrnn:
-        print('Using QRNNs...')
+def read_wiki_articles(filename):
+    articles = []
+    with open(filename, encoding='utf8') as f:
+        lines = f.readlines()
+    current_article = ''
+    for i,line in enumerate(lines):
+        current_article += line
+        if i < len(lines)-2 and lines[i+1] == ' \n' and istitle(lines[i+2]):
+            articles.append(current_article)
+            current_article = ''
+    articles.append(current_article)
+    print(f"Wiki text was split to {len(articles)} articles")
+    return pd.DataFrame({'texts':np.array(articles)})
 
-    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.'
+@dataclass
+class LMHyperParams:
+    dataset_path: str # data_dir
 
-    if subword:
-        # apply sentencepiece tokenization
-        trn_path = dir_path / f'{lang}.wiki.train.tokens'
-        val_path = dir_path / f'{lang}.wiki.valid.tokens'
+    base_lm_path: str = None
+    bidir: bool =False
+    qrnn: bool = True
+    max_vocab: int = 60000
+    tokenizer: Tokenizers = Tokenizers.MOSES
+    pretrained_model: str = None
 
-        read_file(trn_path, 'train')
-        read_file(val_path, 'valid')
-
-        sp = get_sentencepiece(dir_path, trn_path, name, vocab_size=max_vocab)
-
-        lm_type = contrib_data.LanguageModelType.BiLM if bidir else  contrib_data.LanguageModelType.FwdLM
-
-        data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp, bs=bs, bptt=bptt, lm_type=lm_type)
-        itos = data_lm.train_ds.vocab.itos
-        stoi = data_lm.train_ds.vocab.stoi
-    else:
-        # 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)}")
-
-        itos_fname = model_dir / f'itos_{name}.pkl'
-        if not itos_fname.exists():
-            # create the vocabulary
-            cnt = Counter(word for sent in trn_tok for word in sent)
-            itos = [o for o,c in cnt.most_common(n=max_vocab)]
-            itos.insert(1, PAD)  #  set pad id to 1 to conform to fast.ai standard
-            assert UNK in itos, f'Unknown words are expected to have been replaced with {UNK} in the data.'
-
-            # save vocabulary
-            print(f"Saving vocabulary as {itos_fname}")
-            results['itos_fname'] = itos_fname
-            with open(itos_fname, 'wb') as f:
-                pickle.dump(itos, f)
-        else:
-            print("Loading itos:", itos_fname)
-            itos = np.load(itos_fname)
-        vocab = Vocab(itos)
-        stoi = vocab.stoi
-
-        trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok])
-        val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_tok])
-
-        lm_type = contrib_data.LanguageModelType.BiLM if bidir else  contrib_data.LanguageModelType.FwdLM
-
-        # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos))
-        data_lm = TextLMDataBunch.from_ids(path=dir_path, vocab=vocab, train_ids=trn_ids,
-                                           valid_ids=val_ids, bs=bs, bptt=bptt,
-                                           lm_type=lm_type
-                                           )
-
-    print('Size of vocabulary:', len(itos))
-    print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)]))
+    emb_sz:int = 400
+    nh: int = None
+    nl: int = 3
 
     # these hyperparameters are for training on ~100M tokens (e.g. WikiText-103)
     # for training on smaller datasets, more dropout is necessary
-    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])
-        drop_mult = 0.1
-    else:
-        emb_sz, nh, nl = 400, 1150, 3
-        # emb_sz, nh, nl = 400, 1150, 3
-        dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15])
-        drop_mult = 0.1
+    dps = (0.25, 0.1, 0.2, 0.02, 0.15) # consider removing dps & clip from the default hyperparams and put them to train
+    clip: float = 0.12
+    bptt: int = 70
 
-    fastai.text.learner.default_dropout['language'] = dps
+    lang: str = 'en'
+    name: str = None
+    cuda_id: InitVar[int] = 0
 
-    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, model_dir=model_dir.name,
-                       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))
+    def __post_init__(self, cuda_id):
+        if not torch.cuda.is_available():
+            print('CUDA not available. Setting device=-1.')
+            cuda_id = -1
+        torch.cuda.set_device(cuda_id)
+        self.dataset_path = Path(self.dataset_path)
+        self.base_lm_path = Path(self.base_lm_path) if self.base_lm_path is not None else None
+        self.tokenizer = Tokenizers(self.tokenizer) if isinstance(self.tokenizer, str) else self.tokenizer
 
-    learn.true_wd = False
-    print("true_wd: ", learn.true_wd)
+        assert self.dataset_path.exists()
+        self.cache_dir = self.dataset_path / 'models' / self.tokenizer_prefix
+        self.model_dir = self.cache_dir / self.model_name
 
-    if bidir:
-        learn.metrics = [accuracy_fwd, accuracy_bwd]
-    else:
-        learn.metrics = [accuracy]
+        self.model_dir.mkdir(exist_ok=True, parents=True)
+        print('Max vocab:', self.max_vocab)
+        print('Cache dir:', self.cache_dir)
+        print('Model dir:', self.model_dir)
+        self.dps = np.array(self.dps)
+        if self.nh is None: self.nh = 1550 if self.qrnn else 1150
+        if self.name is None: self.name = self.lang
 
-    try:
-        learn.load(f'{model_name}_{name}')
-        print("Weights loaded")
-    except FileNotFoundError:
-        print("Starting from random weights")
-        pass
+    @property
+    def tokenizer_prefix(self): return f"{self.tokenizer.value}{self.max_vocab // 1000}k"
 
-    learn.fit_one_cycle(num_epochs, 5e-3, (0.8, 0.7), wd=1e-7)
+    @property
+    def model_prefix(self): return ('bi' if self.bidir else '') + ('qrnn' if self.qrnn else 'lstm')
 
-    if not subword and max_vocab is None:
+    @property
+    def model_name(self): return f"{self.model_prefix}_{self.name}.m"
+
+    @property
+    def pretrained_fnames(self): return [self.base_lm_path / 'lm_best', self.base_lm_path / '../itos'] if self.base_lm_path else None
+
+    @property
+    def lm_type(self):
+        return contrib_data.LanguageModelType.BiLM if self.bidir else contrib_data.LanguageModelType.FwdLM
+
+    def tokenzier_to_fastai_args(self, trn_data_loading_func, add_moses):
+        tok_func = MosesTokenizerFunc if add_moses else BaseTokenizer
+        if self.tokenizer is Tokenizers.SUBWORD:
+            if self.base_lm_path: # ensure we are using the same sentence piece model
+                shutil.copy(self.base_lm_path / '..' / 'itos.pkl', self.cache_dir)
+                shutil.copy(self.base_lm_path / '..' / 'spm.model', self.cache_dir)
+                shutil.copy(self.base_lm_path / '..' / 'spm.vocab', self.cache_dir)
+            args = get_sentencepiece(self.cache_dir,
+                                     trn_data_loading_func,
+                                     vocab_size=self.max_vocab,
+                                     use_moses=add_moses,
+                                     lang=self.lang)
+
+        elif self.tokenizer is Tokenizers.MOSES:
+            args = dict(tokenizer=Tokenizer(tok_func=tok_func, lang=self.lang, pre_rules=[replace_std_toks], post_rules=[]))
+        elif self.tokenizer is Tokenizers.MOSES_FA:
+            args = dict(tokenizer=Tokenizer(tok_func=tok_func, lang=self.lang))  # use default pre/post rules
+        elif self.tokenizer is Tokenizers.FASTAI:
+            args = dict()
+        else:
+            raise ValueError(
+                f"self.tokenizer has wrong value {self.tokenizer}, Allowed values are taken from {Tokenizers}")
+        return args
+
+    def save_info(self):
+        from dataclasses import asdict
+        vals = {k: (str(v) if isinstance(v, Path) else v) for k,v in asdict(self).items()}
+        vals.pop('name', None)
+        vals.pop('lang', None)
+        vals['tokenizer'] = self.tokenizer.value
+        with (self.model_dir / 'info.json').open("w") as fp: json.dump(vals, fp)
+        print("Saving info", self.model_dir / 'info.json')
+
+    def train_lm(self, num_epochs=20, data_lm=None, bs=70, true_wd=False, drop_mult=0.0, lr=5e-3):
+        data_lm = self.load_wiki_data(bs=bs) if data_lm is None else data_lm
+        learn = self.create_lm_learner(data_lm, drop_mult=drop_mult)
+
+        learn.true_wd = true_wd
+        if num_epochs > 0:
+            if self.pretrained_fnames or self.pretrained_model:
+                print("Training lm from: ", self.pretrained_fnames or self.pretrained_model)
+                if learn.true_wd:
+                    learn.freeze_to(-1)
+                    learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7))
+                    learn.unfreeze()
+                    learn.fit_one_cycle(num_epochs, 1e-3, moms=(0.8, 0.7))
+                else:
+                    learn.freeze_to(-1)
+                    learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7), wd=1e-7)  # TODO Fix the learning rates
+                    learn.unfreeze()
+                    learn.fit_one_cycle(num_epochs, 1e-3, moms=(0.8, 0.7), wd=1e-7)
+            else:
+                print("Training lm from random weights")
+                learn.unfreeze()
+                if not learn.true_wd: learn.fit_one_cycle(num_epochs, lr, (0.8, 0.7), wd=1e-7)
+                else:                 learn.fit_one_cycle(num_epochs, lr, (0.8, 0.7)) # TODO find proper values
+        learn.save("lm_best_with_opt", with_opt=False)
+        learn.save_encoder(ENC_BEST)
+        learn.save(LM_BEST, with_opt=False)
+        print(learn.path)
+
+        self.save_info()
+        return learn
+
+    def create_lm_learner(self, data_lm, dps=None, **kwargs):
+        fastai.text.learner.default_dropout['language'] = dps or self.dps
+        lm_learner = bilm_learner if self.bidir else language_model_learner
+
+        trn_args = dict(tie_weights=True, clip=self.clip, bptt=self.bptt,
+                        pretrained_fnames=self.pretrained_fnames,
+                        pretrained_model=self.pretrained_model)
+        trn_args.update(kwargs)
+        print ("Training args: ", trn_args, "dps: ", dps or self.dps)
+        learn = lm_learner(data_lm, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, pad_token=PAD_TOKEN_ID,
+                           bias=True, qrnn=self.qrnn, model_dir=self.model_dir.relative_to(data_lm.path), **trn_args)
+        # compared to standard Adam, we set beta_1 to 0.8
+        learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99))
+        learn.metrics = [accuracy_fwd, accuracy_bwd] if self.bidir else [accuracy]
+        learn.callback_fns += [partial(CSVLogger, filename=f"{learn.model_dir}/lm-history"),
+                               partial(SaveModelCallback, every='epoch', name='lm')]
+        return learn
+
+    def load_train_text(self):
+        trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens'
+        with open(trn_path) as f:
+            return [line.rstrip('\n') for line in f]
+
+    def load_wiki_data(self, bs=70):
+        trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens'
+        val_path = self.dataset_path / f'{self.lang}.wiki.valid.tokens'
+        tst_path = self.dataset_path / f'{self.lang}.wiki.test.tokens'
+        for path_ in [trn_path, val_path, tst_path]:
+            assert path_.exists(), f'Error: {path_} does not exist.'
+
+        args = self.tokenzier_to_fastai_args(trn_data_loading_func=self.load_train_text, add_moses=False)
+        try:
+            data_lm = TextLMDataBunch.load(self.cache_dir, '.', lm_type=self.lm_type, bs=bs)
+            print("Tokenized data loaded")
+        except FileNotFoundError:
+            print("Running tokenization")
+            data_lm = TextLMDataBunch.from_df(path=self.cache_dir, train_df=read_wiki_articles(trn_path),
+                                              valid_df=read_wiki_articles(val_path),
+                                              classes=None, lm_type=self.lm_type, max_vocab=self.max_vocab,
+                                              bs=bs, text_cols='texts', **args)
+            data_lm.save('.')
+
+        itos, stoi, trn_path = data_lm.vocab.itos, data_lm.vocab.stoi, data_lm.path
+        print('Size of vocabulary:', len(itos))
+        print('First 20 words in vocab:', data_lm.vocab.itos[:20])
+        return data_lm
+
+    @classmethod
+    def from_lm(cls, dataset_path, base_lm_path, **kwargs) -> 'LMHyperParams':
+        base_lm_path = Path(base_lm_path).resolve()
+        dataset_path = Path(dataset_path).resolve()
+        with open(base_lm_path/'info.json', 'r') as f: d = json.load(f)
+        d['dataset_path'] = dataset_path
+        d['base_lm_path'] = base_lm_path
+        d.pop('bs', None)
+        d.pop('drop_mult', None)
+        subword = d.pop('subword', False)
+        tokenizer = d.pop('tokenizer', None)
+        if tokenizer is not None:
+            d['tokenizer'] = Tokenizers(tokenizer)
+        elif subword:
+            d['tokenizer'] = Tokenizers.SUBWORD
+        else:
+            d['tokenizer'] = Tokenizers.MOSES
+
+        d.update(kwargs)
+        return cls(**d)
+
+def validate_lm(self):
+    if not self.exp.subword and self.exp.max_vocab is None:
+        raise NotImplementedError("figure out how to validate and save results")
         # 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(learn.model, tst_ids, bptt)
+        logloss, perplexity = validate(learn.model, tst_ids, self.exp.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}')
-
-    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)
+    fire.Fire(LMHyperParams)

ulmfit/train_clas.py

@@ -2,193 +2,205 @@
 Train a classifier on top of a language model trained with `pretrain_lm.py`.
 Optionally fine-tune LM before.
 """
-import numpy as np
-import pickle
+from sacremoses import MosesTokenizer
 
+import fastai
 import torch
-from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner
-from fastai import fit_one_cycle, accuracy
+
+from fastai import *
+from fastai.callbacks import CSVLogger, SaveModelCallback
+from fastai.text import *
+from fastai_contrib import utils
+
 from fastai_contrib.data import LanguageModelType
 from fastai_contrib.learner import bilm_text_classifier_learner, bilm_learner, accuracy_fwd, accuracy_bwd
-from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists
+from fastai_contrib.utils import PAD, UNK, read_clas_data, PAD_TOKEN_ID, DATASETS, TRN, VAL, TST, ensure_paths_exists, \
+    get_sentencepiece, MosesTokenizerFunc
 from fastai.text.transform import Vocab
 
 import fire
 from collections import Counter
 from pathlib import Path
 
+from ulmfit.pretrain_lm import LMHyperParams, Tokenizers, ENC_BEST
 
-def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_dir='models',
-                   qrnn=False, num_lm_epochs=10,
-                   fine_tune=True, max_vocab=60000, bs=20, bptt=70, name='imdb-clas',
-                   dataset='imdb', bidir=False, ds_pct=1.0, train=True):
-    """
-    :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)
+class CLSHyperParams(LMHyperParams):
+    # dir_path -> data/imdb/
+    use_test_for_validation=False
 
-    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.'
+    bicls_head:str = 'BiPoolingLinearClassifier'
 
-    data_dir = Path(data_dir)
-    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)
+    def __post_init__(self, *args, **kwargs):
+        super().__post_init__(*args, **kwargs)
+        self.dataset_dir=self.dataset_path
+
+    @property
+    def need_fine_tune_lm(self): return not (self.model_dir/f"enc_best.pth").exists()
 
 
-    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}')
+    def train_cls(self, num_lm_epochs, unfreeze=True, bs=40, true_wd=True, drop_mul_lm=0.3, drop_mul_cls=0.5,
+                  use_test_for_validation=False):
+        data_clas, data_lm = self.load_cls_data(bs, use_test_for_validation=use_test_for_validation)
 
-    ensure_paths_exists(data_dir,
-                        dataset_dir,
-                        model_dir,
-                        model_dir/f"{pretrained_fname[0]}.pth",
-                        model_dir/f"{pretrained_fname[1]}.pkl")
-
-    if bidir:
-        print("BiLM")
-        classifier_learner = bilm_text_classifier_learner
-        lm_learner = bilm_learner
-    else:
-        classifier_learner = text_classifier_learner
-        lm_learner = language_model_learner
-
-    lm_type = LanguageModelType.BiLM if bidir else LanguageModelType.FwdLM
-    data_clas, data_lm = get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_type=lm_type)
-
-    if qrnn:
-        emb_sz, nh, nl = 400, 1550, 3
-    else:
-        emb_sz, nh, nl = 400, 1150, 3
-
-    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...', lm_enc_finetuned)
-        learn = lm_learner(
-            data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, qrnn=qrnn,
-            pad_token=PAD_TOKEN_ID,
-            pretrained_fnames=pretrained_fname,
-            path=model_dir.parent, model_dir=model_dir.name,
-            drop_mult=0.3)
-        if bidir:
-            learn.metrics = [accuracy_fwd, accuracy_bwd]
+        if self.need_fine_tune_lm: self.train_lm(num_lm_epochs, data_lm=data_lm, true_wd=true_wd, drop_mult=drop_mul_lm)
+        learn = self.create_cls_learner(data_clas, drop_mult=drop_mul_cls)
+        try:
+            learn.load('cls_last')
+            print("Loading last classifier")
+        except FileNotFoundError:
+            learn.load_encoder(ENC_BEST)
+        if true_wd:
+            learn.true_wd = True
+            print("Starting classifier training")
+            learn.freeze_to(-1)
+            learn.fit_one_cycle(1, 2e-2, moms=(0.8, 0.7))
+            if unfreeze:
+                learn.freeze_to(-2)
+                learn.fit_one_cycle(1, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7))
+                learn.freeze_to(-3)
+                learn.fit_one_cycle(1, slice(5e-3 / (2.6 ** 4), 5e-3), moms=(0.8, 0.7))
+                learn.unfreeze()
+                learn.fit_one_cycle(2, slice(1e-3 / (2.6 ** 4), 1e-3), moms=(0.8, 0.7))
         else:
-            learn.metrics = [accuracy]
-
-        learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7))
-        learn.unfreeze()
-        if num_lm_epochs > 0: learn.fit_one_cycle(num_lm_epochs, 1e-3, moms=(0.8, 0.7))
-
-        # save encoder
-        learn.save_encoder(lm_enc_finetuned)
-
-
-    learn = 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, drop_mult=0.5)
-
-    try:
-        print(f"Loading classifier {model_name}_{name}")
-        learn.load(f'{model_name}_{name}')
-
-    except FileNotFoundError:
-        learn.load_encoder(lm_enc_finetuned)
-        print("loading encoder")
-        train = True
-
-    if train:
-        learn.true_wd = False
-        print("Starting classifier training")
-        learn.fit_one_cycle(1, 5e-2, moms=(0.8, 0.7), wd=1e-7)
-
-        learn.freeze_to(-2)
-        learn.fit_one_cycle(1, slice(5e-2 / (2.6 ** 4), 5e-2), moms=(0.8, 0.7), wd=1e-7)
-
-        learn.freeze_to(-3)
-        learn.fit_one_cycle(1, slice(5e-4 / (2.6 ** 4), 5e-4), moms=(0.8, 0.7), wd=1e-7)
-
-        learn.unfreeze()
-        learn.fit_one_cycle(2, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
-
+            learn.true_wd = False
+            print("Starting classifier training")
+            learn.fit_one_cycle(1, 5e-2, moms=(0.8, 0.7), wd=1e-7)
+            if unfreeze:
+                learn.freeze_to(-2)
+                learn.fit_one_cycle(1, slice(5e-2 / (2.6 ** 4), 5e-2), moms=(0.8, 0.7), wd=1e-7)
+                learn.freeze_to(-3)
+                learn.fit_one_cycle(1, slice(5e-4 / (2.6 ** 4), 5e-4), moms=(0.8, 0.7), wd=1e-7)
+                learn.unfreeze()
+                learn.fit_one_cycle(2, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
         print(f"Saving models at {learn.path / learn.model_dir}")
-        learn.save(f'{model_name}_{name}')
+        learn.save('cls_last', with_opt=False)
+        self.validate_cls('cls_last', bs=bs)
+        self.validate_cls('cls_best', bs=bs)
+        return learn
 
-    results['accuracy'] = learn.recorder.metrics[-1][0]
-    return results
+    def validate_cls(self, save_name='cls_last', bs=40):
+        data_clas, data_lm = self.load_cls_data(bs, use_test_for_validation=True)
+        learn = self.create_cls_learner(data_clas, drop_mult=0.1)
+        learn.load(save_name)
+        print(f"Loss and accuracy using ({save_name}):", learn.validate())
+
+    def create_cls_learner(self, data_clas, dps=None, **kwargs):
+        fastai.text.learner.default_dropout['language'] = dps or self.dps
+        trn_args=dict(bptt=self.bptt, clip=self.clip,)
+        trn_args.update(kwargs)
+        classifier_learner = text_classifier_learner
+        if self.bidir:
+            classifier_learner = bilm_text_classifier_learner
+            trn_args['bicls_head'] = self.bicls_head
+        learn = classifier_learner(data_clas,  pad_token=PAD_TOKEN_ID,
+            path=self.model_dir.parent, model_dir=self.model_dir.name,
+            qrnn=self.qrnn, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, **trn_args)
+        learn.callback_fns += [partial(CSVLogger, filename=f"{learn.model_dir}/cls-history"),
+                               partial(SaveModelCallback, every='improvement', name='cls_best')]
+        return learn
+
+    def load_cls_data(self, bs, **kwargs):
+        if 'imdb' in self.dataset_dir.name:
+            return self.load_cls_data_imdb(bs, **kwargs)
+        else:
+            assert self.tokenizer is Tokenizers.MOSES, "XNLI does not support other tokenizers than Moses"
+            return self.load_cls_data_old_for_xnli(bs, **kwargs)
+
+    def load_cls_data_imdb(self, bs, force=False, use_test_for_validation=False):
+        trn_df = pd.read_csv(self.dataset_path / 'train.csv', header=None)
+        tst_df = pd.read_csv(self.dataset_path / 'test.csv', header=None)
+        unsp_df = pd.read_csv(self.dataset_path / 'unsup.csv', header=None)
+
+        lm_trn_df = pd.concat([unsp_df, trn_df, tst_df])
+        val_len = max(int(len(lm_trn_df) * 0.1), 2)
+        lm_trn_df = lm_trn_df[val_len:]
+        lm_val_df = lm_trn_df[:val_len]
+
+        if use_test_for_validation:
+            val_df = tst_df
+            cls_cache = 'notst'
+        else:
+            val_len = max(int(len(trn_df) * 0.1), 2)
+            trn_len = len(trn_df) - val_len
+            trn_df, val_df = trn_df[:trn_len], trn_df[trn_len:]
+            cls_cache = '.'
+
+        args = self.tokenzier_to_fastai_args(trn_data_loading_func=lambda: trn_df[1], add_moses=True)
+
+        try:
+            if force: raise FileNotFoundError("Forcing reloading of caches")
+            data_lm = TextLMDataBunch.load(self.cache_dir, 'lm', lm_type=self.lm_type, bs=bs)
+            print(f"Tokenized data loaded, lm.trn {len(data_lm.train_ds)}, lm.val {len(data_lm.valid_ds)}")
+        except FileNotFoundError:
+            print(f"Running tokenization...")
+            data_lm = TextLMDataBunch.from_df(path=self.cache_dir, train_df=lm_trn_df, valid_df=lm_val_df,
+                                              max_vocab=self.max_vocab, bs=bs, lm_type=self.lm_type, **args)
+            print(f"Saving tokenized: cls.trn {len(data_lm.train_ds)}, cls.val {len(data_lm.valid_ds)}")
+            data_lm.save('lm')
+
+        try:
+            if force: raise FileNotFoundError("Forcing reloading of caches")
+            data_cls = TextClasDataBunch.load(self.cache_dir, cls_cache, bs=bs)
+            print(f"Tokenized data loaded, cls.trn {len(data_cls.train_ds)}, cls.val {len(data_cls.valid_ds)}")
+        except FileNotFoundError:
+            args['vocab'] = data_lm.vocab  # make sure we use the same vocab for classifcation
+            print(f"Running tokenization...")
+            data_cls = TextClasDataBunch.from_df(path=self.cache_dir, train_df=trn_df, valid_df=val_df,
+                                                 test_df=tst_df, max_vocab=self.max_vocab, bs=bs, **args)
+            print(f"Saving tokenized: cls.trn {len(data_cls.train_ds)}, cls.val {len(data_cls.valid_ds)}")
+            data_cls.save(cls_cache)
+        print('Size of vocabulary:', len(data_lm.vocab.itos))
+        print('First 20 words in vocab:', data_lm.vocab.itos[:20])
+        return data_cls, data_lm
 
 
-def get_datasets(dataset, dataset_dir, bptt, bs, lang, max_vocab, ds_pct, lm_type):
-    tmp_dir = dataset_dir / 'tmp'
-    tmp_dir.mkdir(exist_ok=True)
-    vocab_file = tmp_dir / f'vocab_{lang}.pkl'
-    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 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)
-        vocab = Vocab(itos)
-        stoi = vocab.stoi
-        with open(vocab_file, 'wb') as f:
-            pickle.dump(vocab, f)
-
-        ids = {}
+    def load_cls_data_old_for_xnli(self, bs):
+        tmp_dir = self.cache_dir
+        tmp_dir.mkdir(exist_ok=True)
+        vocab_file = tmp_dir / f'vocab_{self.lang}.pkl'
+        if not (tmp_dir / f'{TRN}_{self.lang}_ids.npy').exists():
+            print('Reading the data...')
+            toks, lbls = read_clas_data(self.dataset_dir, self.dataset_dir.name, self.lang)
+            # create the vocabulary
+            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=self.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)
+            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}_{self.lang}_ids.npy', ids[split])
+                np.save(tmp_dir / f'{split}_{self.lang}_lbl.npy', lbls[split])
+        else:
+            print('Loading the pickled data...')
+            ids, lbls = {}, {}
+            for split in [TRN, VAL, TST]:
+                ids[split] = np.load(tmp_dir / f'{split}_{self.lang}_ids.npy')
+                lbls[split] = np.load(tmp_dir / f'{split}_{self.lang}_lbl.npy')
+            with open(vocab_file, 'rb') as f:
+                vocab = pickle.load(f)
+        print(f'Train size: {len(ids[TRN])}. Valid size: {len(ids[VAL])}. '
+              f'Test size: {len(ids[TST])}.')
         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:
-        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)
-    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"Making the dataset smaller {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[TRN]})
-
-    print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
-    return data_clas, data_lm
+            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=self.bptt, lm_type=self.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[TRN]})
 
+        print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
+        return data_clas, data_lm
 
 if __name__ == '__main__':
-    fire.Fire(new_train_clas)
+    fire.Fire(CLSHyperParams)
+
+##
+