mirror of
https://github.com/wassname/multifit.git
synced 2026-09-09 11:27:26 +08:00
sentencepiece for pretraining
modified: fastai_contrib/utils.py modified: ulmfit/pretrain_lm.py
This commit is contained in:
+29
-16
@@ -14,8 +14,8 @@ 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
|
||||
|
||||
import shutil
|
||||
import pathlib
|
||||
import tarfile
|
||||
from sklearn import model_selection
|
||||
@@ -41,33 +41,38 @@ 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'):
|
||||
def __init__(self, model_dir:PathOrStr):
|
||||
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'))
|
||||
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, dataset:TextDataset, rules:ListRules=None,
|
||||
cache_name:str='tmp', vocab_size:int=30000,
|
||||
model_type:str='unigram', input_sentence_size:int=1E7,
|
||||
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):
|
||||
try:
|
||||
import sentencepiece as spm
|
||||
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)
|
||||
path = pathlib.Path(path)
|
||||
os.makedirs(path / 'models', exist_ok=True)
|
||||
rules = rules if rules else default_rules
|
||||
|
||||
cache_name = 'tmp'
|
||||
|
||||
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))
|
||||
# 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'):
|
||||
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)
|
||||
@@ -75,23 +80,31 @@ def get_sentencepiece(path:PathOrStr, dataset:TextDataset, rules:ListRules=None,
|
||||
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'--model_prefix={path / 'models' / 'spm'} ' \
|
||||
f'--vocab_size={vocab_size} --model_type={model_type} '
|
||||
spm.SentencePieceTrainer.Train(sp_params)
|
||||
|
||||
with open(path / cache_name / 'm.vocab', 'r') as f:
|
||||
with open(path / 'models' / '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 / 'itos.pkl', 'wb'))
|
||||
pickle.dump(vocab, open(path / 'models'/ f'itos_{name}.pkl', 'wb'))
|
||||
|
||||
vocab = Vocab(pickle.load(open(path / 'itos.pkl', 'rb')))
|
||||
spt = SentencepieceTokenizer(path, cache_name)
|
||||
vocab = Vocab(pickle.load(open(path / 'models'/ f'itos_{name}.pkl', 'rb')))
|
||||
spt = SentencepieceTokenizer(path)
|
||||
tokenizer = Tokenizer(tok_func=lambda lang: spt, rules=rules)
|
||||
|
||||
clear_cache_directory(path, cache_name)
|
||||
|
||||
return {'tokenizer': tokenizer, 'vocab': vocab}
|
||||
|
||||
|
||||
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):
|
||||
|
||||
+31
-27
@@ -12,7 +12,7 @@ from fastai import *
|
||||
from fastai.text import *
|
||||
import torch
|
||||
from fastai_contrib.utils import read_file, read_whitespace_file,\
|
||||
DataStump, validate, PAD, UNK
|
||||
DataStump, validate, PAD, UNK, get_sentencepiece
|
||||
|
||||
import pickle
|
||||
|
||||
@@ -24,24 +24,21 @@ from collections import Counter
|
||||
# conda install -c pytorch -c fastai fastai pytorch-nightly [cuda92]
|
||||
# cupy needs to be installed for QRNN
|
||||
|
||||
|
||||
def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab=60000,
|
||||
bs=70, bptt=70, name='wt-103', num_epochs=10):
|
||||
def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vocab=60000,
|
||||
bs=70, bptt=70, name='wt-103', model_dir='models', num_epochs=10):
|
||||
"""
|
||||
:param dir_path: The path to the directory that contains wiki text
|
||||
: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 qrrn: Use a QRNN. Requires installing cupy.
|
||||
:param clean: Train on the clean
|
||||
: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
|
||||
"""
|
||||
|
||||
model_dir = 'models' # removed from params, as it is absolute models location in train_clas and here it is relative
|
||||
if not torch.cuda.is_available():
|
||||
print('CUDA not available. Setting device=-1.')
|
||||
cuda_id = -1
|
||||
@@ -57,13 +54,26 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
|
||||
if qrnn:
|
||||
print('Using QRNNs...')
|
||||
|
||||
trn_path = dir_path / f'{lang}.wiki.train.tokens'
|
||||
val_path = dir_path / f'{lang}.wiki.valid.tokens'
|
||||
tst_path = dir_path / f'{lang}.wiki.test.tokens'
|
||||
trn_path = dir_path / f'{lang}.wiki.train.tokens.unk'
|
||||
val_path = dir_path / f'{lang}.wiki.valid.tokens.unk'
|
||||
tst_path = dir_path / f'{lang}.wiki.test.tokens.unk'
|
||||
for path_ in [trn_path, val_path, tst_path]:
|
||||
assert path_.exists(), f'Error: {path_} does not exist.'
|
||||
|
||||
if clean:
|
||||
if subword:
|
||||
# apply sentencepiece tokenization
|
||||
trn_path = dir_path / f'{lang}.wiki.train.tokens'
|
||||
val_path = dir_path / f'{lang}.wiki.valid.tokens'
|
||||
|
||||
read_file(trn_path, 'train')
|
||||
read_file(val_path, 'valid')
|
||||
|
||||
sp = get_sentencepiece(dir_path, trn_path, name)
|
||||
|
||||
data_lm = TextLMDataBunch.from_csv(dir_path, **sp)
|
||||
itos = data_lm.train_ds.vocab.itos
|
||||
stoi = data_lm.train_ds.vocab.stoi
|
||||
else:
|
||||
# read the already whitespace separated data without any preprocessing
|
||||
trn_tok = read_whitespace_file(trn_path)
|
||||
val_tok = read_whitespace_file(val_path)
|
||||
@@ -77,20 +87,20 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
|
||||
|
||||
vocab = Vocab(itos)
|
||||
stoi = vocab.stoi
|
||||
|
||||
# save vocabulary
|
||||
print(f"Saving vocabulary as {dir_path / model_dir}")
|
||||
with open(dir_path / model_dir / f'itos_{name}.pkl', 'wb') as f:
|
||||
pickle.dump(itos, f)
|
||||
|
||||
|
||||
trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok])
|
||||
val_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in val_tok])
|
||||
|
||||
# data_lm = TextLMDataBunch.from_ids(dir_path, trn_ids, [], val_ids, [], len(itos))
|
||||
data_lm = TextLMDataBunch.from_ids(path=dir_path, vocab=vocab, train_ids=trn_ids,
|
||||
valid_ids=val_ids, bs=bs, bptt=bptt)
|
||||
else:
|
||||
# apply fastai preprocessing and tokenization
|
||||
read_file(trn_path, 'train')
|
||||
read_file(val_path, 'valid')
|
||||
|
||||
data_lm = TextLMDataBunch.from_csv(dir_path, max_vocab=max_vocab)
|
||||
itos = data_lm.train_ds.vocab.itos
|
||||
stoi = data_lm.train_ds.vocab.stoi
|
||||
|
||||
print('Size of vocabulary:', len(itos))
|
||||
print('First 10 words in vocab:', ', '.join([itos[i] for i in range(10)]))
|
||||
@@ -116,17 +126,11 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, clean=True, max_vocab
|
||||
learn.opt_fn = partial(optim.Adam, betas=(0.8, 0.99))
|
||||
learn.true_wd = False
|
||||
|
||||
# save vocabulary
|
||||
print(f"Saving vocabulary as {dir_path / model_dir}")
|
||||
with open(dir_path / model_dir / f'itos_{name}.pkl', 'wb') as f:
|
||||
pickle.dump(itos, f)
|
||||
|
||||
fit_one_cycle(learn, num_epochs, 5e-3, (0.8, 0.7), wd=1e-7)
|
||||
|
||||
if clean and max_vocab is None:
|
||||
if not subword and max_vocab is None:
|
||||
# only if we use the unpreprocessed version and the full vocabulary
|
||||
# are the perplexity results comparable to previous work
|
||||
|
||||
print(f"Validating model performance with test tokens from: {trn_path}")
|
||||
tst_tok = read_whitespace_file(trn_path)
|
||||
tst_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in tst_tok])
|
||||
|
||||
Reference in New Issue
Block a user