diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 5f40940..5bc7b7e 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -8,6 +8,11 @@ import torch from tqdm import tqdm import re import csv +from functools import reduce +from fastai.text.data import TextDataset +from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab, default_rules +from fastai.torch_core import * +from pathlib import Path EOS = '' UNK = '' @@ -17,6 +22,57 @@ PAD_TOKEN_ID = 1 number_match_re = re.compile(r'^([0-9]+[,.]?)+$') number_split_re = re.compile(r'([,.])') +class SentencepieceTokenizer(BaseTokenizer): + def __init__(self, path:PathOrStr, cache_name:str='tmp'): + try: + import sentencepiece as spm + except ImportError: + raise Exception('sentencepiece module is missing: run `pip install sentencepiece`') + self.tok = spm.SentencePieceProcessor() + self.tok.Load(str(Path(path) / cache_name / 'm.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, dataset:TextDataset, rules:ListRules=None, + cache_name:str='tmp', vocab_size:int=30000, + model_type:str='unigram', input_sentence_size:int=1E7, + pad_idx:int=PAD_TOKEN_ID): + try: + import sentencepiece as spm + except ImportError: + raise Exception('sentencepiece module is missing: run `pip install sentencepiece`') + + path = Path(path) + os.makedirs(path / cache_name, exist_ok=True) + rules = rules if rules else default_rules + + if not os.path.isfile(path / cache_name / 'm.model') or not os.path.isfile(path / 'itos.pkl'): + raw_text = reduce(lambda t, rule: rule(t), rules, '\n'.join(dataset.x)) + 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 / cache_name / "m"} ' \ + f'--vocab_size={vocab_size} --model_type={model_type} ' + spm.SentencePieceTrainer.Train(sp_params) + + with open(path / cache_name / 'm.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 / 'itos.pkl', 'wb')) + + vocab = Vocab(pickle.load(open(path / 'itos.pkl', 'rb'))) + spt = SentencepieceTokenizer(path, cache_name) + tokenizer = Tokenizer(tok_func=lambda lang: spt, rules=rules) + + return {'tokenizer': tokenizer, 'vocab': vocab} def replace_number(token): """Replaces a number and returns a list of one or multiple tokens."""