diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 927426a..977b6bb 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -1,18 +1,9 @@ """ Utility methods for data processing. """ -import pandas as pd -import numpy as np import fire -from fastai import F, to_device -import torch -from tqdm import tqdm -import re -import csv - -from functools import reduce -from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab -from fastai.torch_core import * +from fastai import * +from fastai.text import * import shutil import pathlib @@ -22,6 +13,7 @@ from sacremoses import MosesTokenizer from typing import Dict, Tuple, List EOS = '' +BOS = '' UNK = '' PAD = '' SEP = '' # special separator token for NLI @@ -39,67 +31,102 @@ CLASSES = ['neg', 'pos', 'unsup'] number_match_re = re.compile(r'^([0-9]+[,.]?)+$') number_split_re = re.compile(r'([,.])') -class SentencepieceTokenizer(BaseTokenizer): - def __init__(self, model_dir:PathOrStr): +class MosesTokenizerFunc(BaseTokenizer): + "Wrapper around a MosesTokenizer to make it a `BaseTokenizer`." + def __init__(self, lang:str): + 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} @@ -107,7 +134,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): @@ -189,48 +215,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. @@ -249,7 +233,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}...') @@ -268,13 +259,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 diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 3b499b8..e1c4ea3 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -23,41 +23,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 @@ -195,10 +170,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' @@ -206,16 +186,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) diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py index 264be7c..22cf91f 100644 --- a/ulmfit/train_clas.py +++ b/ulmfit/train_clas.py @@ -139,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: