diff --git a/fastai_contrib/data.py b/fastai_contrib/data.py
new file mode 100644
index 0000000..274d1f6
--- /dev/null
+++ b/fastai_contrib/data.py
@@ -0,0 +1,73 @@
+"NLP data loading pipeline. Supports csv, folders, and preprocessed data."
+from fastai.text import *
+from fastai.torch_core import *
+from fastai.text.transform import *
+from fastai.basic_data import *
+from fastai.data_block import *
+
+#region Modified fastai classes
+
+LanguageModelType=Enum('LanguageModelType', 'FwdLM BwdLM BiLM')
+
+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
+        self.first,self.i,self.iter = True,0,0
+        self.n = len(np.concatenate(dataset.x.items)) // self.bs
+        self.max_len,self.num_workers = max_len,0
+
+    def __iter__(self):
+        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([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):
+            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.
+                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)
+            pos += seq_len
+            itr += 1
+            yield res
+
+    def __len__(self) -> int: return int(math.ceil((self.n-1) / self.bptt)) # so that it is always at least 1
+    def __getattr__(self,k:str)->Any: return getattr(self.dataset, k)
+
+    @property
+    def batch_size(self):
+        return self.bs
+
+    @batch_size.setter
+    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)
+        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)
+        return x,y
+
+#endregion
+#region Replaces fastai classes
+
+import fastai.text.data
+fastai.text.data.LanguageModelLoader = LanguageModelLoader # Replace original LanguageModelLoader with new verion
+
+#endregion
\ No newline at end of file
diff --git a/fastai_contrib/learner.py b/fastai_contrib/learner.py
new file mode 100644
index 0000000..e060b7a
--- /dev/null
+++ b/fastai_contrib/learner.py
@@ -0,0 +1,101 @@
+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 *
+
+#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,
+                  pretrained_fnames:OptStrTuple=None, **kwargs) -> 'LanguageLearner':
+    "Create a `Learner` with a language model."
+    dps = default_dropout['language'] * drop_mult
+    vocab_size = len(data.vocab.itos)
+    model = get_bilm(vocab_size, emb_sz, nh, nl, pad_token, input_p=dps[0], output_p=dps[1],
+                weight_p=dps[2], embed_p=dps[3], hidden_p=dps[4], tie_weights=tie_weights, bias=bias, qrnn=qrnn)
+    learn = LanguageLearner(data, model, bptt, split_func=bilm_split, **kwargs)
+    if pretrained_model is not None:
+        model_path = untar_data(pretrained_model, data=False)
+        fnames = [list(model_path.glob(f'*.{ext}'))[0] for ext in ['pth', 'pkl']]
+        learn.load_pretrained(*fnames)
+        learn.freeze()
+    if pretrained_fnames is not None:
+        fnames = [learn.path/learn.model_dir/f'{fn}.{ext}' for fn,ext in zip(pretrained_fnames, ['pth', 'pkl'])]
+        learn.load_pretrained(*fnames)
+        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."
+    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.')
+    else:
+        return convert_weights_with_prefix(wgts, stoi_wgts, itos_new, prefix='')
+
+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
+    return wgts
+
+#endregion
+#region Replace code in fastai
+
+import fastai.text.learner
+fastai.text.learner.convert_weights = convert_weights
+
+#endregion
\ No newline at end of file
diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py
new file mode 100644
index 0000000..57ccc47
--- /dev/null
+++ b/fastai_contrib/models.py
@@ -0,0 +1,152 @@
+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):
+        super().__init__()
+        self.fwd_lm = fwd_lm
+        self.bwd_lm = bwd_lm
+
+    def __getitem__(self, idx):
+        return BiLMModel(self.fwd_lm[idx], self.bwd_lm[idx])
+
+    def __len__(self):
+        return len(self.fwd_lm)
+
+    def stack(self, fwd_o, bwd_o):
+        if is_listy(fwd_o):
+            return [self.stack(f, b) for f,b in zip(fwd_o,bwd_o)]
+        else:
+            return torch.stack([fwd_o, bwd_o], dim=len(fwd_o.shape))
+
+    def forward(self, input):
+        if len(input.shape) == 3: # sl, bs, tracks
+            f = input[..., 0]
+            b = input[..., 1]
+        elif len(input.shape) == 2: # sl, bs - support during classification mode
+            f = input
+            b = torch.flip(input, [0])
+        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)
+
+    def reset(self):
+        "Reset the hidden states of underlaying lms."
+        self.fwd_lm.reset()
+        self.bwd_lm.reset()
+
+
+class BiPoolingLinearClassifier(nn.Module):
+    "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
+        elif len(output.size()) == 4:
+            sl, bs, 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 = self.layers(x)
+            return x, raw_outputs, outputs
+
+
+class AvgPoolingLinearClassifier(nn.Module):
+    "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
+        elif len(output.size()) == 4:
+            sl, bs, em_sz, passes = output.size()
+
+            avgpool = (self.pool(output[...,0], bs, False) + self.pool(output[..., 1], bs, False))/2
+            mxpool = (self.pool(output[...,0], bs, True) +self.pool(output[..., 1], bs, True))/2
+            x = torch.cat([(output[-1][..., 0]+output[-1][..., 1])/2, mxpool, avgpool], 1)
+            x = self.layers(x)
+            return x, raw_outputs, outputs
+
+
+def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, tie_weights:bool=True,
+                       qrnn:bool=False, bias:bool=True, bidir:bool=False, output_p:float=0.4, hidden_p:float=0.2, input_p:float=0.6,
+                       embed_p:float=0.1, weight_p:float=0.5)->nn.Module:
+    "Create a two AWD-LSTM one for each direction "
+    fwd_rnn_enc = RNNCore(vocab_sz, emb_sz, n_hid=n_hid, n_layers=n_layers, pad_token=pad_token, qrnn=qrnn, bidir=bidir,
+                          hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
+    bwd_rnn_enc = RNNCore(vocab_sz, emb_sz, n_hid=n_hid, n_layers=n_layers, pad_token=pad_token, qrnn=qrnn, bidir=bidir,
+                          hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
+    enc = None
+    if tie_weights:
+        enc = fwd_rnn_enc.encoder
+        fwd_rnn_enc.encoder.weight = enc.weight
+        bwd_rnn_enc.encoder.weight = enc.weight
+
+    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)))
+
+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:
+    "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))
+    model.reset()
+    return model
+
+#endregion
\ No newline at end of file
diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py
index 3017ffc..2dd21f9 100644
--- a/fastai_contrib/utils.py
+++ b/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
     """
@@ -225,7 +225,7 @@ def read_imdb(dir_path, lang, split, spm_path=None) -> Tuple[List[List[str]], Li
         reader = csv.reader(f)
         for row in reader:
             label, text = row
-            lbls.append(label)
+            lbls.append(int(label))
             raw_tokens = mt.tokenize(text, return_str=True).split(' ')
             
             tokens = []
@@ -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:
diff --git a/results/bilcls_qrnn-wt-103-wd0 b/results/bilcls_qrnn-wt-103-wd0
new file mode 100644
index 0000000..4f2a68f
--- /dev/null
+++ b/results/bilcls_qrnn-wt-103-wd0
@@ -0,0 +1,26 @@
+pretraining
+0 - 8 lost, about 0.30 0.32 after 5 epochs
+8      3.744213    4.021358    0.308377      0.325672
+9      3.700674    4.021499    0.308637      0.325881
+10     3.674045    4.022058    0.308656      0.325903
+--- crash---
+$ python -m ulmfit.train_clas --data_dir data --model_dir data/wiki/wikitext-103/models --pretrain_name=bilm-wt-103 --qrnn=True --name 'concat-2x' --cuda-id=0 --bs 40 --train=True --bidir=True
+Dataset: imdb. Language: en.
+Using QRNNs...
+BiLM
+Loading the pickled data...
+Train size: 22500. Valid size: 2500. Test size: 25000.
+loading encoder
+Starting classifier training
+epoch  train_loss  valid_loss  accuracy
+1      0.346874    0.276838    0.881200
+epoch  train_loss  valid_loss  accuracy
+1      0.321121    0.243946    0.901200
+epoch  train_loss  valid_loss  accuracy
+1      0.303174    0.234627    0.908400
+epoch  train_loss  valid_loss  accuracy
+1      0.295207    0.227533    0.912800
+2      0.269754    0.221328    0.914400
+Saving models at data/wiki/wikitext-103/models
+accuracy: tensor(0.9144)
+
diff --git a/results/bilm_qrnn-wt-103-unk-wd1 b/results/bilm_qrnn-wt-103-unk-wd1
new file mode 100644
index 0000000..dc5c8ce
--- /dev/null
+++ b/results/bilm_qrnn-wt-103-unk-wd1
@@ -0,0 +1,20 @@
+(fastaiv1) n-waves@GV100:~/workspace/ulmfit-multilingual$ python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --cuda-id=0 -
+-name=bilm-wt-103-unk --num_epochs=10  --bs=64  --bidir=True                                                                               
+Batch size: 64                                                                                                                             
+Max vocab: 60000                                                                                                                           
+Using QRNNs...                                                                                                                             
+Loading itos: data/wiki/wikitext-103-unk/models/itos_bilm-wt-103-unk.pkl                                                                   
+Size of vocabulary: 60001                                                                                                                  
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a                                                                       
+Starting from random weights                                                                                                               
+epoch  train_loss  valid_loss  accuracy_fwd  accuracy_bwd                                                                                  
+1      4.184118    4.113817    0.308417      0.318290                                                                                      
+2      4.106002    4.039978    0.312220      0.320315                                                                                      
+3      4.160282    4.086650    0.306327      0.314152                                                                                      
+4      4.136789    4.049762    0.309202      0.317373                                                                                      
+5      4.086520    4.001534    0.313158      0.321543                                                                                      
+6      4.036415    3.960494    0.318333      0.325623                                                                                      
+7      3.988382    3.913568    0.324409      0.330398                                                                                      
+8      3.937409    3.874283    0.328386      0.334786                                                                                      
+9      3.912215    3.856325    0.330593      0.337102
+10     3.876535    3.848783    0.331202      0.337649
\ No newline at end of file
diff --git a/results/bilm_qrnn-wt-103-wd0 b/results/bilm_qrnn-wt-103-wd0
new file mode 100644
index 0000000..291173e
--- /dev/null
+++ b/results/bilm_qrnn-wt-103-wd0
@@ -0,0 +1,20 @@
+(fastaiv1) n-waves@GV100:~/workspace/ulmfit-multilingual$ python -m ulmfit.pretrain_lm
+ki/wikitext-103 --qrnn=True --cuda-id=0 --name=bilm-wt-103 --num_epochs=10  --bs=64  --bidir=True
+Batch size: 64
+Max vocab: 60000
+Using QRNNs...
+Loading itos: data/wiki/wikitext-103/models/itos_bilm-wt-103.pkl
+Size of vocabulary: 60001
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a
+Starting from random weights
+epoch  train_loss  valid_loss  accuracy_fwd  accuracy_bwd
+1      4.141750    4.114931    0.307700      0.320260
+2      4.078711    4.061186    0.310221      0.320550
+3      4.108722    4.081286    0.308383      0.316431
+4      4.054437    4.040497    0.311229      0.319658
+5      4.014880    3.966010    0.318588      0.327607
+6      3.926770    3.897305    0.326874      0.333171
+7      3.855364    3.816452    0.335632      0.342045
+8      3.761600    3.745798    0.343413      0.350675
+
+10     3.614097    3.687730    0.351651      0.358280
\ No newline at end of file
diff --git a/results/lm_lstm-wt-103-wd0 b/results/lm_lstm-wt-103-wd0
new file mode 100644
index 0000000..41c178f
--- /dev/null
+++ b/results/lm_lstm-wt-103-wd0
@@ -0,0 +1,22 @@
+(fastaiv1) pczapla@galatea ~/w/ulmfit-multilingual ❯❯❯ time python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=False --cuda-id=1 --name=wt-103 --num_epochs=10  --bs=32
+Batch size: 32
+Max vocab: 60000
+Loading itos: data/wiki/wikitext-103/models/itos_wt-103.pkl
+Size of vocabulary: 60001
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a
+Starting from random weights
+epoch  train_loss  valid_loss  accuracy
+1      4.317837    4.245714    0.309034
+2      4.347099    4.251623    0.306932
+3      4.360198    4.302850    0.302014
+4      4.329998    4.264225    0.306828
+5      4.301852    4.209807    0.311763
+6      4.218213    4.131647    0.319867
+7      4.167365    4.047796    0.327259
+8      4.095695    3.980298    0.336255
+9      4.032371    3.919622    0.343671
+10     3.983160    3.906227    0.345863
+Saving models at data/wiki/wikitext-103/models
+Saving optimiser state at data/wiki/wikitext-103/models/lstm3_wt-103_state.pth
+accuracy: tensor(0.3461)
+python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=False --cuda-id=1   44147.64s user 16746.69s system 99% cpu 16:58:47.57 total
\ No newline at end of file
diff --git a/results/lm_qrnn-wt-103-unk-wd0 b/results/lm_qrnn-wt-103-unk-wd0
new file mode 100644
index 0000000..d009387
--- /dev/null
+++ b/results/lm_qrnn-wt-103-unk-wd0
@@ -0,0 +1,23 @@
+time python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --cuda-id=0 --name=wt-103 --num_epochs=10  --bs=32
+Batch size: 32
+Max vocab: 60000
+Using QRNNs...
+Loading itos: data/wiki/wikitext-103-unk/models/itos_wt-103.pkl
+Size of vocabulary: 60001
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a
+Starting from random weights
+epoch  train_loss  valid_loss  accuracy
+1      4.393613    4.282537    0.305135
+2      4.394575    4.299973    0.300028
+3      4.444056    4.336682    0.295496
+4      4.414292    4.332908    0.297408
+5      4.406248    4.269964    0.302962
+6      4.327769    4.209948    0.309025
+7      4.244819    4.140769    0.315675
+8      4.210000    4.075574    0.323544
+9      4.139524    4.034881    0.328550
+10     4.150556    4.021361    0.331256
+Saving models at data/wiki/wikitext-103-unk/models
+Saving optimiser state at data/wiki/wikitext-103-unk/models/qrnn3_wt-103_state.pth
+accuracy: tensor(0.3312)
+python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True    --bs=3  32119.72s user 12439.18s system 99% cpu 12:25:03.65 total
\ No newline at end of file
diff --git a/results/lm_qrnn-wt-103-unk-wd1 b/results/lm_qrnn-wt-103-unk-wd1
new file mode 100644
index 0000000..0ee4a69
--- /dev/null
+++ b/results/lm_qrnn-wt-103-unk-wd1
@@ -0,0 +1,24 @@
+ime python -m ulmfit.pretrain_lm data/wiki/wikitext-103-unk --qrnn=True --cuda-id=0 --name=wt-103-wd1 --num_epochs=10  --bs=32
+Batch size: 32
+Max vocab: 60000
+Using QRNNs...
+Saving vocabulary as data/wiki/wikitext-103-unk/models
+Size of vocabulary: 60001
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a
+true_wd:  True
+Starting from random weights
+epoch  train_loss  valid_loss  accuracy
+1      4.355600    4.271729    0.306827
+2      4.384312    4.288582    0.299482
+3      4.527577    4.366407    0.289836
+4      4.519645    4.387311    0.288076
+5      4.514816    4.371242    0.289862
+6      4.506230    4.337718    0.294208
+7      4.444272    4.295349    0.298786
+8      4.432104    4.257432    0.301785
+9      4.416849    4.242803    0.304147
+10     4.389035    4.239752    0.304871
+Saving models at data/wiki/wikitext-103-unk/models
+Saving optimiser state at data/wiki/wikitext-103-unk/models/qrnn3_wt-103-wd1_state.pth
+itos_fname: data/wiki/wikitext-103-unk/models/itos_wt-103-wd1.pkl
+accuracy:   tensor(0.2829)
\ No newline at end of file
diff --git a/results/lm_qrnn-wt-103-wd0 b/results/lm_qrnn-wt-103-wd0
new file mode 100644
index 0000000..33d243b
--- /dev/null
+++ b/results/lm_qrnn-wt-103-wd0
@@ -0,0 +1,20 @@
+time python -m ulmfit.pretrain_lm data/wiki/wikitext-103 --qrnn=True --cuda-id=1 --name=wt-103 --num_epochs=10  --bs=32                             ✘ 1
+Batch size: 32
+Max vocab: 60000
+Using QRNNs...
+Size of vocabulary: 60001
+First 10 words in vocab: the, <pad>, ,, ., of, and, to, in, <eos>, a
+Saving vocabulary as data/wiki/wikitext-103/models
+epoch  train_loss  valid_loss  accuracy
+1      4.328219    4.268609    0.306414
+2      4.351323    4.282290    0.300560
+3      4.376134    4.345460    0.294145
+4      4.394314    4.319661    0.298142
+5      4.337362    4.265630    0.303437
+6      4.295763    4.199119    0.309354
+7      4.165526    4.121037    0.318062
+8      4.151789    4.060387    0.325899
+9      4.066542    4.014568    0.332732
+10     4.013142    3.997691    0.335299
+Saving models at data/wiki/wikitext-103/models
+Saving optimiser state at data/wiki/wikitext-103/models/qrnn3_wt-103_state.pth
\ No newline at end of file
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
index f17bbfb..9b80a6b 100644
--- a/tests/test_end_to_end.py
+++ b/tests/test_end_to_end.py
@@ -1,52 +1,53 @@
 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'
+import fastai.core
+fastai.core.turn_off_parallel_execution=True
 
-    # 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 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 check_data_exists():
+def get_test_data():
     data = get_data_folder()
+    wt = data / "wiki" / "wikitext-2"
+    imdb = data / "imdb"
 
-    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
+    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=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(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,31 +57,29 @@ 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)
-
-    delete_test_models()
+        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,
+        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')
 
 
 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(
@@ -89,8 +88,8 @@ def test_ulmfit_sentencepiece_end_to_end():
         cuda_id=cuda_id,
         qrnn=True,
         subword=True,
-        max_vocab=1000,
-        bs=80,
+        max_vocab=100,
+        bs=2,
         num_epochs=1,
         name=lm_name,
     )
@@ -100,4 +99,7 @@ def test_ulmfit_sentencepiece_end_to_end():
     # NOTE: ds_pct is not available for sentencepiece -- tests are on the complete dataset
     #       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
+
diff --git a/tests/test_text_data.py b/tests/test_text_data.py
new file mode 100644
index 0000000..6b3dabf
--- /dev/null
+++ b/tests/test_text_data.py
@@ -0,0 +1,59 @@
+import pytest
+import fastai.text
+
+from fastai import *
+from fastai.text import *
+
+import fastai_contrib.data as contrib_data
+
+def text_df(labels):
+    data = []
+    texts = ["fast ai is a cool project", "hello world"] * 20
+    for ind, text in enumerate(texts):
+        sample = {}
+        sample["label"] = labels[ind%len(labels)]
+        sample["text"] = text
+        data.append(sample)
+    return pd.DataFrame(data)
+
+###################### UPDATED CODE
+def test_should_load_backwards_lm():
+    path = untar_data(URLs.IMDB_SAMPLE)
+    df = text_df(['neg','pos'])
+
+    data = TextLMDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=0, text_cols=["text"], bs=2,
+                                   lm_type=contrib_data.LanguageModelType.BwdLM,
+                                   ld_cls=contrib_data.LanguageModelLoader)
+    lml = data.train_dl.dl
+    lml.data = lml.batchify(np.concatenate([lml.dataset.x.items[i] for i in range(len(lml.dataset))]))
+    batch = lml.get_batch(lml.data, 0, 70)
+
+    assert batch[0].shape == (70, lml.bs)
+    assert batch[1].shape == (70*lml.bs,)
+
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:,0]]
+    np.testing.assert_array_equal(as_text[:5], ["world", "hello", '1', 'xxfld', 'project',])
+
+def test_should_load_bi_lm():
+    path = untar_data(URLs.IMDB_SAMPLE)
+    df = text_df(['neg', 'pos'])
+
+    data = TextLMDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=0, text_cols=["text"], bs=2,
+                                   lm_type=contrib_data.LanguageModelType.BiLM,
+                                   ld_cls=contrib_data.LanguageModelLoader)
+    lml = data.train_dl.dl
+    lml.data = lml.batchify(np.concatenate([lml.dataset.x.items[i] for i in range(len(lml.dataset))]))
+    batch = lml.get_batch(lml.data, 0, 70)
+
+    assert batch[0].shape == (70, lml.bs, 2)
+    assert batch[1].shape == (70*lml.bs, 2)
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:, 0, 0]]
+    np.testing.assert_array_equal(as_text[:7], "xxfld 1 fast ai is a cool".split())
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:,0,1]]
+    np.testing.assert_array_equal(as_text[:5], ["world", "hello", '1', 'xxfld', 'project',])
+
+###################### NEW CODE
+
diff --git a/tests/test_text_train.py b/tests/test_text_train.py
new file mode 100644
index 0000000..0103939
--- /dev/null
+++ b/tests/test_text_train.py
@@ -0,0 +1,101 @@
+import pytest
+from fastai import *
+from fastai.text import *
+
+pytestmark = pytest.mark.integration
+
+print(sys.path)
+import fastai_contrib.data as contrib_data
+
+from fastai_contrib.learner import bilm_learner, accuracy_fwd, bilm_text_classifier_learner
+
+
+def read_file(fname):
+    texts = []
+    with open(fname, 'r') as f:
+        texts = f.readlines()
+    labels = [0] * len(texts)
+    df = pd.DataFrame({'labels':labels, 'texts':texts}, columns = ['labels', 'texts'])
+    return df
+
+def prep_human_numbers():
+    path = untar_data(URLs.HUMAN_NUMBERS)
+    df_trn = read_file(path/'train.txt')
+    df_val = read_file(path/'valid.txt')
+    return path, df_trn, df_val
+
+def manual_seed(seed=42):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+@pytest.fixture(scope="module")
+def learn():
+    path, df_trn, df_val = prep_human_numbers()
+    data = TextLMDataBunch.from_df(path, df_trn, df_val, tokenizer=Tokenizer(BaseTokenizer))
+    learn = language_model_learner(data, emb_sz=100, nl=1, drop_mult=0.1)
+    learn.fit_one_cycle(4, 5e-3)
+    return learn
+
+def text_df(n_labels):
+    data = []
+    texts = ["fast ai is a cool project", "hello world"]
+    for ind, text in enumerate(texts):
+        sample = {}
+        for label in range(n_labels): sample[label] = ind%2
+        sample["text"] = text
+        data.append(sample)
+    df = pd.DataFrame(data)
+    return df
+
+###################### NEW CODE
+
+def test_val_loss(learn):
+    assert learn.validate()[1] > 0.5
+
+
+def test_bilm_classifier_loads_encoder():
+    n_labels=2
+    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data', 'tmp')
+    os.makedirs(path)
+    try:
+        df = text_df(n_labels=1)
+        lmdf = df#[["text"]]
+        print(lmdf.head())
+        lmdata = TextLMDataBunch.from_df(path, lmdf, lmdf, tokenizer=Tokenizer(BaseTokenizer),
+                                       lm_type=contrib_data.LanguageModelType.BiLM)
+        learn = bilm_learner(lmdata, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False)
+        learn.save_encoder("enc")
+        data = TextClasDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=list(range(n_labels)), text_cols=["text"])
+        classifier = bilm_text_classifier_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False)
+        print(last_layer(classifier.model), )
+        classifier.load_encoder("enc")
+        classifier.fit(1)
+    finally:
+        shutil.rmtree(path)
+
+
+def test_bilm_lstm_can_be_trained():
+    manual_seed()
+    path, df_trn, df_val = prep_human_numbers()
+    data = TextLMDataBunch.from_df(path, df_trn, df_val, tokenizer=Tokenizer(BaseTokenizer),
+                                   lm_type = contrib_data.LanguageModelType.BiLM)
+
+    learn = bilm_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False)
+    learn.metrics = [accuracy_fwd]
+    learn.fit_one_cycle(2, 5e-3)
+    assert learn.validate()[1] > 0.3
+
+
+def test_bwdlm_lstm_can_be_trained():
+    manual_seed()
+    path, df_trn, df_val = prep_human_numbers()
+    data = TextLMDataBunch.from_df(path, df_trn, df_val, tokenizer=Tokenizer(BaseTokenizer),
+                                   lm_type = contrib_data.LanguageModelType.BwdLM)
+
+    learn = language_model_learner(data, emb_sz=100, nl=1, drop_mult=0.1, qrnn=False)
+    learn.fit_one_cycle(2, 5e-3)
+    assert learn.validate()[1] > 0.3
diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py
index cca6a35..942475a 100644
--- a/ulmfit/create_wikitext.py
+++ b/ulmfit/create_wikitext.py
@@ -100,7 +100,7 @@ def main(args):
     write_wikitext(lrg_wiki_train, text_iter, mt, 98000000, mode='a')
     all_wiki_train = all_wiki / f'{args.lang}.wiki.train.tokens'
     copyfile(lrg_wiki_train, all_wiki_train)
-    write_wikitext(lrg_wiki_train, text_iter, mt,  None, mode='a') # TODO fix it (change lrg to all)
+    write_wikitext(all_wiki_train, text_iter, mt,  None, mode='a')
 
 if __name__ == '__main__':
 
diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py
index 64b2da2..aaca9e7 100644
--- a/ulmfit/pretrain_lm.py
+++ b/ulmfit/pretrain_lm.py
@@ -6,26 +6,27 @@ to be split.
 """
 import fastai
 import fire
-import numpy as np
 
 from fastai import *
 from fastai.text import *
 import torch
-from fastai_contrib.utils import read_file, read_whitespace_file,\
-    DataStump, validate, PAD, UNK, get_sentencepiece
-
+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
 import pickle
 
 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
 
+
 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, ds_pct=1.0):
+                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
@@ -38,9 +39,10 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
     :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 = {}
-    model_dir = 'models' # removed from params, as it is absolute models location in train_clas and here it is relative
+
     if not torch.cuda.is_available():
         print('CUDA not available. Setting device=-1.')
         cuda_id = -1
@@ -48,7 +50,7 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
 
     dir_path = Path(dir_path)
     assert dir_path.exists()
-    model_dir = Path(model_dir)
+    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)
@@ -72,7 +74,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
 
         sp = get_sentencepiece(dir_path, trn_path, name, vocab_size=max_vocab)
 
-        data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp)
+        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:
@@ -84,28 +88,35 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
             val_tok = val_tok[:max(20, int(len(val_tok) * ds_pct))]
             print(f"Limiting data sets to {ds_pct*100}%, trn {len(trn_tok)}, val: {len(val_tok)}")
 
-        # create the vocabulary
-        cnt = Counter(word for sent in trn_tok for word in sent)
-        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.'
+        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
 
-        # save vocabulary
-        itos_fname = model_dir / f'itos_{name}.pkl'
-        print(f"Saving vocabulary as {itos_fname}")
-        results['itos_fname'] = itos_fname
-        with open(itos_fname, 'wb') as f:
-            pickle.dump(itos, f)
-
         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)
+                                           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)]))
@@ -123,15 +134,31 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
         dps = np.array([0.25, 0.1, 0.2, 0.02, 0.15])
         drop_mult = 0.1
 
-    fastai.text.learner.default_dropout['language'] = dps * drop_mult
-    learn = language_model_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,
-                           bias=True, qrnn=qrnn, clip=0.12)
+    fastai.text.learner.default_dropout['language'] = dps
+
+    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))
-    learn.true_wd = False
 
-    fit_one_cycle(learn, num_epochs, 5e-3, (0.8, 0.7), wd=1e-7)
+    learn.true_wd = False
+    print("true_wd: ", learn.true_wd)
+
+    if bidir:
+        learn.metrics = [accuracy_fwd, accuracy_bwd]
+    else:
+        learn.metrics = [accuracy]
+
+    try:
+        learn.load(f'{model_name}_{name}')
+        print("Weights loaded")
+    except FileNotFoundError:
+        print("Starting from random weights")
+        pass
+
+    learn.fit_one_cycle(num_epochs, 5e-3, (0.8, 0.7), wd=1e-7)
 
     if not subword and max_vocab is None:
         # only if we use the unpreprocessed version and the full vocabulary
diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py
index 4bfbb8e..7b88963 100644
--- a/ulmfit/train_clas.py
+++ b/ulmfit/train_clas.py
@@ -7,7 +7,9 @@ import pickle
 
 import torch
 from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner
-from fastai import fit_one_cycle
+from fastai import fit_one_cycle, accuracy
+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.text.transform import Vocab
 
@@ -17,9 +19,9 @@ from pathlib import Path
 
 
 def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_dir='models',
-                   qrnn=False,
-                   fine_tune=True, max_vocab=30000, bs=20, bptt=70, name='imdb-clas',
-                   dataset='imdb', ds_pct=1.0):
+                   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
@@ -49,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)
@@ -67,17 +69,87 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
                         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]
+        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)
+
+        print(f"Saving models at {learn.path / learn.model_dir}")
+        learn.save(f'{model_name}_{name}')
+
+    results['accuracy'] = learn.recorder.metrics[-1][0]
+    return results
+
+
+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] 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)
@@ -100,63 +172,23 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
             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"Makeing the dataset smaller {ds_pct}")
-        for split in [TRN, VAL, TST]:
-            ids[split] = ids[split][:int(len(ids[split])*ds_pct)]
-
-
-    data_lm = TextLMDataBunch.from_ids(path=tmp_dir, vocab=vocab, train_ids=ids[TRN],
-                                       valid_ids=ids[VAL], bs=bs, bptt=bptt)
-
-    # TODO TextClasDataBunch allows tst_ids as input, but not tst_lbls?
+        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)
+        train_lbls=lbls[TRN], valid_lbls=lbls[VAL], bs=bs, classes={l:l for l in lbls[TRN]})
 
-    if qrnn:
-        emb_sz, nh, nl = 400, 1550, 3
-    else:
-        emb_sz, nh, nl = 400, 1150, 3
-    learn = language_model_learner(
-        data_lm, bptt=bptt, emb_sz=emb_sz, nh=nh, nl=nl, qrnn=qrnn,
-        pad_token=PAD_TOKEN_ID,
-        pretrained_fnames=pretrained_fname,
-        path=model_dir.parent, model_dir=model_dir.name)
-    lm_enc_finetuned  = f"{lm_name}_{dataset}_enc"
-    if fine_tune and not (model_dir / f"lm_enc_finetuned.pth").exists():
-        print('Fine-tuning the language model...')
-        learn.unfreeze()
-        learn.fit(2, slice(1e-4, 1e-2))
+    print(f"Sizes of train_ds {len(data_clas.train_ds)}, valid_ds {len(data_clas.valid_ds)}")
+    return data_clas, data_lm
 
-        # save encoder
-        learn.save_encoder(lm_enc_finetuned)
-
-    print("Starting classifier training")
-    learn = text_classifier_learner(data_clas, bptt=bptt, pad_token=PAD_TOKEN_ID,
-                                  path=model_dir.parent, model_dir=model_dir.name,
-                                  qrnn=qrnn, emb_sz=emb_sz, nh=nh, nl=nl)
-
-    learn.load_encoder(lm_enc_finetuned)
-
-    learn.fit_one_cycle(1, 2e-2, moms=(0.8, 0.7), wd=1e-7)
-
-    learn.freeze_to(-2)
-    learn.fit_one_cycle(1, slice(1e-2 / (2.6 ** 4), 1e-2), moms=(0.8, 0.7), wd=1e-7)
-
-    learn.freeze_to(-3)
-    learn.fit_one_cycle(1, slice(5e-3 / (2.6 ** 4), 5e-3), moms=(0.8, 0.7), wd=1e-7)
-
-    learn.unfreeze()
-    learn.fit_one_cycle(2, slice(1e-3 / (2.6 ** 4), 1e-3), moms=(0.8, 0.7), wd=1e-7)
-    results['accuracy'] = learn.validate()[1]
-    print(f"Saving models at {learn.path / learn.model_dir}")
-    learn.save(f'{model_name}_{name}')
-    return results
 
 if __name__ == '__main__':
     fire.Fire(new_train_clas)