From 3e8f9f8b5f057123918dfdef7eeb55f8320ccd7a Mon Sep 17 00:00:00 2001 From: Aayush Date: Thu, 15 Nov 2018 23:01:49 +0530 Subject: [PATCH 1/6] [WIP] Sub-word tokenization with sentencepiece For resolving issue #3. @eisenjulian let's use this branch. I've written some skeleton code (untested currently) to be used around your tokenizer. You can commit that into fastai_contrib for our purpose. --- ulmfit/create_wikitext.py | 102 ++++++++++++++++++++++++++++++++++++-- 1 file changed, 97 insertions(+), 5 deletions(-) diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py index 5d24284..b6f3dd9 100644 --- a/ulmfit/create_wikitext.py +++ b/ulmfit/create_wikitext.py @@ -11,6 +11,7 @@ import json from shutil import copyfile from sacremoses import MosesTokenizer +from fastai_contrib.utils import get_sentencepiece, replace_number, UNK def get_texts(root): @@ -27,10 +28,12 @@ def get_texts(root): yield text -def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'): +def write_wikitext(file_path, text_iter, tok, num_tokens, mode='w'): + total_num_tokens = 0 print(f'Writing to {file_path}...') i = 0 + with open(file_path, mode, encoding='utf-8') as f_out: for i, text in enumerate(text_iter): @@ -39,7 +42,7 @@ def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'): paragraphs = text.split('\n') for paragraph in paragraphs: - tokenized = mt.tokenize(paragraph.strip(), return_str=True) + tokenized = tok.tokenize(paragraph.strip(), return_str=True) tokenized_paragraphs.append(tokenized) tokens = tokenized.split(' ') # split on whitespace to keep newlines @@ -64,6 +67,70 @@ def write_wikitext(file_path, text_iter, mt, num_tokens, mode='w'): print('{}. # documents: {:,}. # tokens: {:,}.'.format( file_path, i, total_num_tokens)) +def build_vocab(file_path, cutoff=3): + counter = Counter() + with open(file_path, 'r', encoding='utf-8') as f: + for i, line in enumerate(f): + tokens = line.strip().split(' ') + [''] + counter.update(tokens) + vocab = {} + in_vocab_count = 0 + OOV_count = 0 + for token, count in counter.most_common(): + if count >= cutoff: + vocab[token] = count + in_vocab_count += count + else: + OOV_count += count + print('OOV ratio: %.4f.' % (OOV_count / (in_vocab_count + OOV_count))) + return vocab + + +def limit_vocab(unk_path, vocab): + """ + https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-limit_vocab-py + :param unk_path: + :param vocab: + :return: + """ + temp_file_path = unk_path.with_name(unk_path.name + '.temp') + total_num_tokens = 0 + print(f'Limiting vocab in {unk_path}. Writing to {unk_path}.') + with open(unk_path, 'r', encoding='utf-8') as f_in, open(temp_file_path, 'w', encoding='utf-8') as f_out: + for line in f_in: + tokens = [x for x in line.strip().split(' ') if x] + tokens = [token if token in vocab else UNK for token in tokens] + # Ensures there's a space between tokens, including the last word, + # newline, and the first word of the next line + tokens = tokens + ['\n'] + total_num_tokens += len(tokens) + tokens = [''] + tokens + line = ' '.join(tokens) + f_out.write(line) + print(f'{unk_path.name}. # of tokens: {total_num_tokens}') + temp_file_path.replace(unk_path) + +def replace_numbers(text_iter, unk_path): + """ + Replace numbers as in Smerity's script: + https://gist.github.com/Smerity/94af5902aa9498817c92d1e71eb2f87b#file-post_process-py + :param file_path: + :param unk_path: + :return: + """ + print(f'Replacing numbers in file. Writing to {unk_path}.') + with open(unk_path, 'w', encoding='utf-8') as f: + for text in text_iter: + raw_tokens = line.strip().split(' ') + tokens = [] + for token in raw_tokens: + tokens.append(replace_number(token)) + # Starting each line with a blank line is required + # Some systems replace \n with and assume, like in PTB, everything is space separated + tokens = [''] + tokens + ['\n'] + line = ' '.join(tokens) + f.write(line) + def main(args): @@ -72,7 +139,13 @@ def main(args): assert input_path.exists(), f'Error: {input_path} does not exist.' output.mkdir(exist_ok=True) - mt = MosesTokenizer(args.lang) + if args.subword: + # TO DO load the text corpus + # TO DO make get_sentencepiece return path to spm model + spm_path = get_sentencepiece(output, corpus) + tok = SentencepieceTokenizer(spm_path) + else: + tok = MosesTokenizer(args.lang) sml_wiki = output / f'{args.lang}-2' lrg_wiki = output / f'{args.lang}-100' @@ -84,17 +157,33 @@ def main(args): splits = ['train', 'valid', 'test'] token_nums = [2000000, 200000, 200000] for split, token_num in zip(splits, token_nums): + # TO DO maybe replace the numbers before tokenizing + unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk' + replace_numbers(text_iter, unk_path) + sml_file_path = sml_wiki / f'{args.lang}.wiki.{split}.tokens' - write_wikitext(sml_file_path, text_iter, mt, token_num) + write_wikitext(sml_file_path, text_iter, tok, token_num) lrg_file_path = lrg_wiki / f'{args.lang}.wiki.{split}.tokens' # copy the content of the small file to the large file print(f'Copying {sml_file_path} to {lrg_file_path}.') copyfile(sml_file_path, lrg_file_path) + sml_vocab = build_vocab(sml_wiki_train) + print(f'{args.lang}-2 vocab size: {len(sml_vocab)}') + lrg_vocab = build_vocab(lrg_wiki_train) + print(f'{args.lang}-100 vocab size: {len(lrg_vocab)}') + # add the new articles to the existing ones lrg_wiki_train = lrg_wiki / f'{args.lang}.wiki.train.tokens' - write_wikitext(lrg_wiki_train, text_iter, mt, 98000000, mode='a') + write_wikitext(lrg_wiki_train, text_iter, tok, 98000000, mode='a') + + # replace words not in the vocab with + if not args.subword: + for wiki, vocab in zip([sml_wiki, lrg_wiki], [sml_vocab, lrg_vocab]): + for split in splits: + unk_path = wiki / f'{args.lang}.wiki.{split}.tokens.unk' + limit_vocab(unk_path, vocab) if __name__ == '__main__': @@ -110,5 +199,8 @@ if __name__ == '__main__': parser.add_argument('-l', '--lang', required=True, help='the iso code of the language of the Wikipedia ' 'documents, e.g. en, fr, de, etc.') + parser.add_argument('-sw', '--subword', default=False, + help='set to use sub-word tokenization (sentencepiece)' + 'default tokenization method is Moses.') args = parser.parse_args() main(args) From 1cc49e666fb3053a5f526f535f5bc94119eadee4 Mon Sep 17 00:00:00 2001 From: Aayush Date: Thu, 15 Nov 2018 23:13:16 +0530 Subject: [PATCH 2/6] Quick fixes to imports etc. --- ulmfit/create_wikitext.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/ulmfit/create_wikitext.py b/ulmfit/create_wikitext.py index b6f3dd9..072b90e 100644 --- a/ulmfit/create_wikitext.py +++ b/ulmfit/create_wikitext.py @@ -6,12 +6,14 @@ Articles are tokenized using the Moses tokenizer. Articles with least than """ import argparse from pathlib import Path +from collections import Counter import json from shutil import copyfile from sacremoses import MosesTokenizer -from fastai_contrib.utils import get_sentencepiece, replace_number, UNK +from fastai_contrib.utils import replace_number, UNK +from fastai_contrib.tokenizers import get_sentencepiece, SentencepieceTokenizer def get_texts(root): From b9c587d9bdce14c30c4ef98fcf15340ed5b3774c Mon Sep 17 00:00:00 2001 From: "NAUSICAA\\Julian" Date: Thu, 15 Nov 2018 19:18:42 -0300 Subject: [PATCH 3/6] Adding SentencePieceTokenizer --- fastai_contrib/utils.py | 56 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) 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.""" From 59c8852b5d5e48686136ddf62b9e69a2934d935c Mon Sep 17 00:00:00 2001 From: aayush Date: Fri, 16 Nov 2018 23:14:24 +0530 Subject: [PATCH 4/6] sentencepiece for pretraining modified: fastai_contrib/utils.py modified: ulmfit/pretrain_lm.py --- fastai_contrib/utils.py | 45 ++++++++++++++++++++------------ ulmfit/pretrain_lm.py | 58 ++++++++++++++++++++++------------------- 2 files changed, 60 insertions(+), 43 deletions(-) diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index abbe436..b60611b 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -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): diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 6370955..12ce63f 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -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]) From e72b19499c0c7a909d7a139259dbb85a97f2c0d8 Mon Sep 17 00:00:00 2001 From: Aayush Date: Fri, 16 Nov 2018 23:54:48 +0530 Subject: [PATCH 5/6] fix indent --- ulmfit/pretrain_lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 12ce63f..4a16530 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -91,7 +91,7 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo # 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) + pickle.dump(itos, f) trn_ids = np.array([([stoi.get(w, stoi[UNK]) for w in s]) for s in trn_tok]) From 36b056a465f02351fd2694945e20b8ad88c0454d Mon Sep 17 00:00:00 2001 From: Piotr Czapla Date: Fri, 16 Nov 2018 23:31:16 +0100 Subject: [PATCH 6/6] Fix issues discovered during execution of end-to-end test. --- fastai_contrib/utils.py | 3 +-- tests/test_end_to_end.py | 5 ++++- ulmfit/pretrain_lm.py | 6 +++--- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py index 77d44a3..9004132 100644 --- a/fastai_contrib/utils.py +++ b/fastai_contrib/utils.py @@ -11,7 +11,6 @@ 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 * @@ -80,7 +79,7 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N 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"--model_prefix={path / 'models' / 'spm'} " \ f'--vocab_size={vocab_size} --model_type={model_type} ' spm.SentencePieceTrainer.Train(sp_params) diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 77c77be..008c488 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -23,11 +23,13 @@ def check_data_exists(): def test_pretrain_lm(): imdb,wt2 = check_data_exists() lm_name="end-to-end-test-quick" + cuda_id=0 results = ulmfit.pretrain_lm.pretrain_lm( dir_path=wt2, lang='en', + cuda_id=cuda_id, qrnn=True, - clean=True, + subword=False, max_vocab=1000, bs=80, num_epochs=1, @@ -40,6 +42,7 @@ def test_pretrain_lm(): data_dir=get_data_folder(), lang='en', pretrain_name=lm_name, model_dir=wt2/'models', qrnn=True, + cuda_id=cuda_id, fine_tune=True, max_vocab=1000, bs=20, bptt=70, name=lm_name+'-imdb-clas', diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py index 26dfa53..0847ca6 100644 --- a/ulmfit/pretrain_lm.py +++ b/ulmfit/pretrain_lm.py @@ -56,9 +56,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo if qrnn: print('Using QRNNs...') - 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' + 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' for path_ in [trn_path, val_path, tst_path]: assert path_.exists(), f'Error: {path_} does not exist.'