Support for running sentencepiece with train_clas

Committing into this branch. Only changes are in `read_xnli` and `read_imdb`.

fastai_contrib/utils.py

@@ -47,8 +47,10 @@ class SentencepieceTokenizer(BaseTokenizer):
             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
 
@@ -77,11 +79,11 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N
         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)} ' \
+        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} '
+                    f"--vocab_size={vocab_size} --model_type={model_type} "
         spm.SentencePieceTrainer.Train(sp_params)
   
         with open(path / 'models' / 'spm.vocab', 'r') as f:
@@ -199,40 +201,54 @@ def prepare_imdb(file_path: str, prepare_lm = False):
         df_val.to_csv(LM_PATH / 'test.csv', header=False, index=False)
 
 
-def read_imdb(dir_path, lang, split) -> Tuple[List[List[str]], List[str]]:
+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(label)
             raw_tokens = mt.tokenize(text, return_str=True).split(' ') + [EOS]
+            
             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)
-            toks.append(tokens)
+            
+            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) -> Tuple[List[List[str]], List[str]]:
+def read_xnli(dir_path, lang, split, spm_path=None) -> Tuple[List[List[str]], List[str]]:
     """
     Reads XNLI data.
     :param dir_path: the path to the xnli folder
     :param lang: the language
     :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 = XNLI_PATHS[split]
@@ -242,6 +258,10 @@ def read_xnli(dir_path, lang, split) -> Tuple[List[List[str]], List[str]]:
         file_name = 'xnli.dev.en.tsv' if split == VAL else 'xnli.test.en.tsv'
         file_path = f'XNLI-MT-1.0/xnli/{file_name}'
     file_path = dir_path / file_path
+    
+    if spm_path is not None:
+        sp = SentencepieceTokenizer(spm_path)
+
     toks, lbls = [], []
     print(f'Reading {file_path}...')
     with open(file_path, encoding='utf-8') as f:
@@ -257,9 +277,15 @@ def read_xnli(dir_path, lang, split) -> Tuple[List[List[str]], List[str]]:
                 if ex_lang != lang:
                     continue
                 premise, hypo, label = row[-3], row[-2], row[1]
+            
             # TODO add BOS
-            premise_toks = premise.split(' ') + [EOS]
-            hypo_toks = hypo.split(' ') + [EOS]
+            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]
+            
             toks.append(premise_toks + [SEP] + hypo_toks)
             lbls.append(label)
     return toks, lbls
@@ -389,4 +415,4 @@ class TextReader():
 
 
 if __name__ == "__main__":
-    fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz
+    fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz