diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 9004132..3017ffc 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -11,7 +11,7 @@ import re import csv from functools import reduce -from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab, default_rules +from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab from fastai.torch_core import * import shutil @@ -47,11 +47,14 @@ 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 + def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=None, vocab_size:int=30000, model_type:str='unigram', input_sentence_size:int=1E7, pad_idx:int=PAD_TOKEN_ID): @@ -61,26 +64,27 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N raise Exception('sentencepiece module is missing: run `pip install sentencepiece`') path = pathlib.Path(path) - os.makedirs(path / 'models', exist_ok=True) - rules = rules if rules else default_rules - 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 [] + # 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 / f'itos_{name}.pkl'): + 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_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)} ' \ + 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: @@ -88,11 +92,12 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N vocab[0] = UNK vocab[pad_idx] = PAD - pickle.dump(vocab, open(path / 'models'/ f'itos_{name}.pkl', 'wb')) + pickle.dump(vocab, open(path / 'models' / f'itos_{name}.pkl', 'wb')) - vocab = Vocab(pickle.load(open(path / 'models'/ f'itos_{name}.pkl', 'rb'))) - spt = SentencepieceTokenizer(path) - tokenizer = Tokenizer(tok_func=lambda lang: spt, rules=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) clear_cache_directory(path, cache_name) @@ -198,40 +203,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] + 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) - 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] @@ -241,6 +260,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: @@ -256,9 +279,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 @@ -388,4 +417,4 @@ class TextReader(): if __name__ == "__main__": - fire.Fire() # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz \ No newline at end of file + fire.Fire() # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index f1e3c74..67d72fb 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -73,7 +73,7 @@ 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, **sp) + data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp) itos = data_lm.train_ds.vocab.itos stoi = data_lm.train_ds.vocab.stoi else: @@ -85,24 +85,26 @@ 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)}") - itos_fn=dir_path / model_dir / f'itos_{name}.pkl' - if not itos_fn.exists(): - itos_fn.parent.mkdir(exist_ok=True) + 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.' + vocab = Vocab(itos) + stoi = vocab.stoi + # save vocabulary - print(f"Saving vocabulary as {dir_path / model_dir}") - results['itos_fname'] = itos_fn - with open(results['itos_fname'], 'wb') as f: + + 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_fn) - itos = np.load(itos_fn) - + print("Loading itos:", itos_fname) + itos = np.load(itos_fname) vocab = Vocab(itos) stoi = vocab.stoi @@ -117,7 +119,6 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo lm_type=lm_type ) - print('Size of vocabulary:', len(itos)) print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)]))