diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 688e28d..1eab00b 100644 --- a/fastai_contrib/utils.py +++ b/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 \ No newline at end of file + fire.Fire() # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz