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

fastai_contrib/utils.py

@@ -12,11 +12,15 @@ from sklearn import model_selection
 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
+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']
@@ -31,11 +35,10 @@ CLASSES = ['neg', 'pos', 'unsup']
 number_match_re = re.compile(r'^([0-9]+[,.]?)+$')
 number_split_re = re.compile(r'([,.])')
 
-# 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):
+        super().__init__(lang=lang)
         self.tok = MosesTokenizer(lang)
 
     def tokenizer(self, t:str) -> List[str]:
@@ -69,9 +72,9 @@ class SentencePieceTokenizer(Tokenizer):
         toks = tok.sp.EncodeAsPieces(" ".join(toks))
         return toks
 
-def get_sentencepiece(cache_dir:PathOrStr, load_text, 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'):
+                      use_moses=False, lang='en'):
     try:
         import sentencepiece as spm
     except ImportError:

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

@@ -15,7 +15,8 @@ from fastai.text import *
 from fastai.callbacks.tracker import SaveModelCallback
 import torch
 from fastai_contrib.utils import read_file, read_whitespace_file, \
-    validate, PAD, UNK, get_sentencepiece, read_clas_data, TRN, VAL, TST, PAD_TOKEN_ID
+    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
 
@@ -35,7 +36,7 @@ class Tokenizers(Enum):
     FASTAI='f'
 
 def istitle(line):
-    return len(re.findall(r'^ = [^=]* = $', line)) != 0
+    return len(re.findall(r'^ ?= [^=]* = ?$', line)) != 0
 
 def read_wiki_articles(filename):
     articles = []
@@ -48,6 +49,7 @@ def read_wiki_articles(filename):
             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)})
 
 @dataclass
@@ -112,6 +114,30 @@ class LMHyperParams:
     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()}
@@ -126,11 +152,6 @@ class LMHyperParams:
         learn = self.create_lm_learner(data_lm, drop_mult=drop_mult)
 
         learn.true_wd = true_wd
-        # try:
-        #     learn.load("lm_best_with_opt")
-        #     print("Continuing training")
-        # except FileNotFoundError:
-        #     pass
         if num_epochs > 0:
             if self.pretrained_fnames or self.pretrained_model:
                 print("Training lm from: ", self.pretrained_fnames or self.pretrained_model)
@@ -186,83 +207,19 @@ class LMHyperParams:
         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.'
-        if self.tokenizer is Tokenizers.SUBWORD:
-            sp = get_sentencepiece(self.cache_dir,
-                                   self.load_train_text,
-                                   self.name,
-                                   vocab_size=self.max_vocab,
-                                   use_moses=False,
-                                   lang=self.lang)
 
-            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('.')
+        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('.')
 
-
-        elif self.tokenizer is Tokenizers.MOSES:
-            # read the already whitespace separated data without any preprocessing
-            trn_tok = read_whitespace_file(trn_path)
-            val_tok = read_whitespace_file(val_path)
-            itos_fname = self.cache_dir / f'itos.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=self.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}")
-                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])
-
-            # data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos))
-            data_lm = TextLMDataBunch.from_ids(path=self.dataset_path, vocab=vocab, train_ids=trn_ids,
-                                               valid_ids=val_ids, bs=bs, bptt=self.bptt,
-                                               lm_type=self.lm_type)
-        elif self.tokenizer is Tokenizers.MOSES_FA:
-
-            try:
-                data_lm = TextLMDataBunch.load(self.cache_dir, '.', lm_type=self.lm_type, bs=bs)
-                print("Tokenized data loaded")
-            except FileNotFoundError:
-                print("Running tokenization")
-
-                # wikitext is pretokenized with Moses
-                pretokenized = Tokenizer(tok_func=BaseTokenizer, lang='en', pre_rules=None, post_rules=None)
-                data_lm = TextLMDataBunch.from_df(path=self.cache_dir, train_df=read_wiki_articles(trn_path),
-                                                  valid_df=read_wiki_articles(val_path), tokenizer=pretokenized,
-                                                  classes=None, lm_type=self.lm_type,
-                                                  max_vocab=self.max_vocab, bs=bs, text_cols='texts')
-                data_lm.save('.')
-        elif self.tokenizer is Tokenizers.FASTAI:
-            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')
-                data_lm.save('.')
-        else:
-            raise ValueError(f"self.tokenizer has wrong value {self.tokenizer}, Allowed values are taken from {Tokenizers}")
         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])

ulmfit/train_clas.py

@@ -10,33 +10,20 @@ import torch
 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, \
-    get_sentencepiece
+    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
 
-
-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 CLSHyperParams(LMHyperParams):
     # dir_path -> data/imdb/
     use_test_for_validation=False
@@ -138,28 +125,7 @@ class CLSHyperParams(LMHyperParams):
             trn_df, val_df = trn_df[:trn_len], trn_df[trn_len:]
             cls_cache = '.'
 
-        if self.tokenizer is Tokenizers.SUBWORD:
-            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:
-            args = dict(tokenizer=Tokenizer(tok_func=MosesTokenizerFunc, lang='en')) # 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}")
+        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")
@@ -235,3 +201,6 @@ class CLSHyperParams(LMHyperParams):
 
 if __name__ == '__main__':
     fire.Fire(CLSHyperParams)
+
+##
+