Merge branch 'refactor' of https://github.com/n-waves/ulmfit-multilingual into polyglot-lm

fastai_contrib/data.py

@@ -13,24 +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 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)
+        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)
@@ -49,17 +49,17 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
     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,13 +1,11 @@
 from torch.nn import CrossEntropyLoss
 
-from fastai.metrics import 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 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,
@@ -92,10 +90,14 @@ def convert_weights_with_prefix(wgts:Weights, stoi_wgts:Dict[str,int], itos_new:
         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()

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,7 +143,8 @@ 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,

fastai_contrib/utils.py

@@ -1,17 +1,9 @@
 """
 Utility methods for data processing.
 """
-import pandas as pd
-import numpy as np
 import fire
-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,6 +13,7 @@ from sacremoses import MosesTokenizer
 from typing import Dict, Tuple, List
 
 EOS = '<eos>'
+BOS = '<bos>'
 UNK = '<unk>'
 PAD = '<pad>'
 SEP = '<sep>'  # special separator token for NLI
@@ -38,67 +31,104 @@ 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):
+# FIXME: coping of tokens from one sentencepiece model to another does not work for 50% of tokens
+# FIXME: tokens in sentencepiece are uppercase eventhough post-transformation will convert them to lowercase
+class MosesTokenizerFunc(BaseTokenizer):
+    "Wrapper around a MosesTokenizer to make it a `BaseTokenizer`."
+    def __init__(self, lang:str):
+        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, 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):
+                      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 []
 
-    if not os.path.isfile(path / 'models' / 'spm.model') or not os.path.isfile(path / 'models' / f'itos_{name}.pkl'):
+    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
+
+    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 = [line.rstrip('\n') for line in open(trn_path)]
-        text = list(filter(None, text))
-        raw_text = reduce(lambda t, rule: rule(t), pre_rules, '\n'.join(text)) # FIXME: possibly does not work with pre_rules
-        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)
+        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))
 
-        with open(path / 'models' / 'spm.vocab', 'r') as f:
+        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}
 
 
@@ -106,7 +136,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):
@@ -188,48 +217,6 @@ def prepare_imdb(file_path: str, prepare_lm = 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)
 
-
-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.
@@ -248,7 +235,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}...')
@@ -267,13 +261,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
@@ -290,7 +280,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]

tests/test_end_to_end.py

@@ -56,7 +56,7 @@ def test_ulmfit_works_with_relative_paths():
     exp = ulmfit.pretrain_lm.LMHyperParams(
         dataset_path=wt2.relative_to(Path.cwd()),
         lang='en',
-        qrnn=True,
+        qrnn=False,
         max_vocab=1000,
         name=lm_name,
         cuda_id=cuda_id)
@@ -83,7 +83,7 @@ def test_ulmfit_default_end_to_end():
     exp = ulmfit.pretrain_lm.LMHyperParams(
         dataset_path=wt2,
         lang='en',
-        qrnn=True,
+        qrnn=False,
         max_vocab=1000,
         name=lm_name,
         cuda_id=cuda_id)
@@ -105,7 +105,7 @@ def test_ulmfit_fastai_end_to_end():
         dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
+        qrnn=False,
         tokenizer='f',
         max_vocab=100,
         name=lm_name,
@@ -124,7 +124,7 @@ def test_ulmfit_fastai_bidir_end_to_end():
         dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
+        qrnn=False,
         bidir=True,
         tokenizer='f',
         max_vocab=100,
@@ -144,7 +144,7 @@ def test_ulmfit_moses_fa_bidir_end_to_end():
         dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
+        qrnn=False,
         bidir=True,
         tokenizer='vf',
         max_vocab=100,
@@ -168,15 +168,15 @@ def test_ulmfit_sentencepiece_end_to_end():
         dataset_path=wt2,
         lang='en',
         cuda_id=cuda_id,
-        qrnn=True,
+        qrnn=False,
         tokenizer=ulmfit.pretrain_lm.Tokenizers.SUBWORD,
-        max_vocab=100,
+        max_vocab=200,
         name=lm_name,
     )
     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, )
+    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__":

ulmfit/pretrain_lm.py

@@ -24,41 +24,16 @@ 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
-
-
-# """
-# :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
-# """
 LM_BEST = "lm_best"
 ENC_BEST = "enc_best"
 
 
 class Tokenizers(Enum):
-    SUBWORD='sb'
+    SUBWORD='sp'
     MOSES='v'
     MOSES_FA='vf'
     FASTAI='f'
 
-# tokenizers ={
-#     Tok.MOSES: MosesTok,
-#     Tok.SUBWORD: SentencepieceTok,
-#     Tok.FASTAI: FastaiTok
-# }
-
 def istitle(line):
     return len(re.findall(r'^ = [^=]* = $', line)) != 0
 
@@ -196,10 +171,15 @@ class LMHyperParams:
         # 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}/cls-history"),
+        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'
@@ -207,16 +187,25 @@ class LMHyperParams:
         for path_ in [trn_path, val_path, tst_path]:
             assert path_.exists(), f'Error: {path_} does not exist.'
         if self.tokenizer is Tokenizers.SUBWORD:
-            # apply sentencepiece tokenization
-            trn_path = self.dataset_path / f'{self.lang}.wiki.train.tokens'
-            val_path = self.dataset_path / f'{self.lang}.wiki.valid.tokens'
+            sp = get_sentencepiece(self.cache_dir,
+                                   self.load_train_text,
+                                   self.name,
+                                   vocab_size=self.max_vocab,
+                                   use_moses=False,
+                                   lang=self.lang)
 
-            read_file(trn_path, 'train')
-            read_file(val_path, 'valid')
+            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, **sp,
+                                                  max_vocab=self.max_vocab, bs=bs, text_cols='texts')
+                data_lm.save('.')
 
-            sp = get_sentencepiece(self.dataset_path, trn_path, self.name, vocab_size=self.max_vocab)
 
-            data_lm = TextLMDataBunch.from_csv(self.dataset_path, 'train.csv', **sp, bs=bs, bptt=self.bptt, lm_type=self.lm_type)
         elif self.tokenizer is Tokenizers.MOSES:
             # read the already whitespace separated data without any preprocessing
             trn_tok = read_whitespace_file(trn_path)

ulmfit/train_clas.py

@@ -5,15 +5,12 @@ Optionally fine-tune LM before.
 from sacremoses import MosesTokenizer
 
 import fastai
-import numpy as np
-import pickle
+import torch
 
 from fastai import *
 from fastai.callbacks import CSVLogger, SaveModelCallback
 from fastai.text import *
 
-import torch
-from fastai.text import TextLMDataBunch, TextClasDataBunch, language_model_learner, text_classifier_learner
 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, \
@@ -90,8 +87,8 @@ class CLSHyperParams(LMHyperParams):
                 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('cls_last', with_opt=False)
-        self.validate_cls('cls_last')
-        self.validate_cls('cls_best')
+        self.validate_cls('cls_last', bs=bs)
+        self.validate_cls('cls_best', bs=bs)
         return learn
 
     def validate_cls(self, save_name='cls_last', bs=40):
@@ -142,11 +139,18 @@ class CLSHyperParams(LMHyperParams):
             cls_cache = '.'
 
         if self.tokenizer is Tokenizers.SUBWORD:
-            args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv',
-                                     self.name, vocab_size=self.max_vocab, pre_rules=[], post_rules=[])
-            if self.tokenizer is Tokenizers.SUBWORD:
-                args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv',
-                                         self.name, vocab_size=self.max_vocab, pre_rules=[], post_rules=[])
+            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,
+                                     lambda: trn_df[1],
+                                     self.name,
+                                     vocab_size=self.max_vocab,
+                                     lang='en',
+                                     use_moses=True)
+
+            # TODO remove migration of tokens for SentencePiece as more than 50% of tokens are different in imdb
         elif self.tokenizer is Tokenizers.MOSES:
             args = dict(tokenizer=Tokenizer(tok_func=MosesTokenizerFunc, lang='en', pre_rules=[], post_rules=[]))
         elif self.tokenizer is Tokenizers.MOSES_FA: