Merge pull request #16 from n-waves/sentencepiece_fixes

Sentencepiece Fixes
This commit is contained in:
Julian Eisenschlos
2018-11-20 17:14:41 -03:00
committed by GitHub
+10 -8
View File
@@ -64,16 +64,17 @@ 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 None
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:
@@ -91,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'), pre_rules=rules, post_rules=[])
clear_cache_directory(path, cache_name)