""" Utility methods for data processing. """ import fire from fastai import * from fastai.text import * import shutil import pathlib import tarfile from sklearn import model_selection from sacremoses import MosesTokenizer from typing import Dict, Tuple, List EOS = 'xxeos' # fastai does not use eos, but we do SEP = 'xxsep' # special separator token for NLI def replace_std_toks(x:str) -> str: "Replace standard token names with fastai supported tokens" # We change tokens to f'xx{token_name}' as it is not split by Moses tokenizer, # while f'<{token_name}>' is being split to: '<' f'{token_name}' '>' return x.replace('', UNK).replace('', BOS).replace('', EOS) PAD_TOKEN_ID = 1 IMDB, XNLI, TRN, VAL, TST, EN = 'imdb', 'xnli', 'train', 'val', 'test', 'en' DATASETS = ['imdb', 'xnli'] XNLI_PATHS = { TRN: 'XNLI-MT-1.0/multinli/multinli.train.%s.tsv', VAL: 'XNLI-1.0/xnli.dev.tsv', TST: 'XNLI-1.0/xnli.test.tsv' } CLASSES = ['neg', 'pos', 'unsup'] number_match_re = re.compile(r'^([0-9]+[,.]?)+$') number_split_re = re.compile(r'([,.])') class MosesTokenizerFunc(BaseTokenizer): "Wrapper around a MosesTokenizer to make it a `BaseTokenizer`." def __init__(self, lang:str): super().__init__(lang=lang) self.tok = MosesTokenizer(lang) def tokenizer(self, t:str) -> List[str]: return self.tok.tokenize(t, return_str=False, escape=False) def add_special_cases(self, toks:Collection[str]): for w in toks: assert len(self.tokenizer(w))==1, f"Tokenizer is unable to keep {w} as one token!" class SentencePieceTokenizer(Tokenizer): "Put together rules and a tokenizer function to tokenize text with multiprocessing." def __init__(self, spm_model, lang:str='en', pre_rules:ListRules=None, post_rules:ListRules=None, special_cases:Collection[str]=None, n_cpus:int=None, use_moses=False): super().__init__(self.tok_fun_with_sp, lang, pre_rules, post_rules, special_cases, n_cpus) self.spm_model = spm_model self.use_moses = use_moses def tok_fun_with_sp(self, lang): try: import sentencepiece as spm except ImportError: raise Exception('sentencepiece module is missing: run `pip install sentencepiece`') tok = MosesTokenizerFunc(lang) if self.use_moses else BaseTokenizer(lang) tok.sp = spm.SentencePieceProcessor() tok.sp.Load(str(self.spm_model)) return tok def process_text(self, t:str, tok:BaseTokenizer) -> List[str]: "Process one text `t` with tokenizer `tok`." toks = super().process_text(t, tok) toks = tok.sp.EncodeAsPieces(" ".join(toks)) return toks def get_sentencepiece(cache_dir:PathOrStr, load_text,pre_rules:ListRules=None, post_rules:ListRules=None, vocab_size:int=30000, model_type:str='unigram', input_sentence_size:int=1E7, use_moses=False, lang='en'): try: import sentencepiece as spm except ImportError: raise Exception('sentencepiece module is missing: run `pip install sentencepiece`') cache_dir = pathlib.Path(cache_dir) pre_rules = pre_rules if pre_rules is not None else defaults.text_pre_rules post_rules = post_rules if post_rules is not None else defaults.text_post_rules special_cases = defaults.text_spec_tok if not os.path.isfile(cache_dir / 'spm.model') or not os.path.isfile(cache_dir / f'itos.pkl'): # load the text from the train tokens file text = load_text() text = filter(lambda x: len(x.rstrip(" ")), text) text = (reduce(lambda t, rule: rule(t), pre_rules, line) for line in text) if use_moses: mt = MosesTokenizer(lang) splitter = lambda t: mt.tokenize(t, return_str=False, escape=False) else: splitter = lambda t: t.split() def cleanup_n_postprocess(t): t = splitter(t) for r in post_rules: t = r(t) return ' '.join(t) text = map(cleanup_n_postprocess, text) raw_text_path = cache_dir / 'all_text.txt' with open(raw_text_path, 'w') as f: f.write("\n".join(text)) sp_params = [ f"--input={raw_text_path}", f"--character_coverage=1.0", f"--unk_id={len(defaults.text_spec_tok)}", f"--pad_id=-1", f"--bos_id=-1", f"--eos_id=-1", f"--max_sentence_length=20480", f"--input_sentence_size={int(input_sentence_size)}", f"--user_defined_symbols={','.join(special_cases)}", f"--model_prefix={cache_dir/'spm'}", f"--vocab_size={vocab_size} --model_type={model_type}"] spm.SentencePieceTrainer.Train(" ".join(sp_params)) with open(cache_dir / 'spm.vocab', 'r') as f: vocab = [line.split('\t')[0] for line in f.readlines()] pickle.dump(vocab, open(cache_dir/ f'itos.pkl', 'wb')) # todo add post rules vocab = Vocab(pickle.load(open(cache_dir / f'itos.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 = SentencePieceTokenizer(cache_dir/'spm.model', use_moses=use_moses, lang=lang, pre_rules=pre_rules, post_rules=post_rules) 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): for fname in (path/label).glob('*.*'): texts.append(fname.open('r', encoding='utf-8').read()) labels.append(idx) return np.array(texts), np.array(labels) def ensure_paths_exists(*paths, message="One or more required files cannot be found."): error = False for path in paths: if not path.exists(): print(f'Error: {path} does not exist.') error = True if error: raise FileNotFoundError(message) def get_data_folder() -> Path: """ return data folder to use for future processing """ return (pathlib.Path(__file__).parent.parent / "data") def get_scripts_folder(): """ return data folder to use for future processing """ return (pathlib.Path(__file__).parent.parent) def prepare_imdb(file_path: str, prepare_lm = False): """ function to extract aclImdb and combine into fastai standard format of labels and then text columns Args: file_path: path to the aclImdb.tgz prepare_lm (bool): prepare file for language model finetuning Returns: None """ file_path = pathlib.Path(file_path) dir_path = pathlib.Path(file_path.parent / 'aclImdb').resolve() assert tarfile.is_tarfile(file_path), "this is not a valid targz file" if not dir_path.exists(): print(f"Extracting {file_path} to {dir_path}. This may take a long time...") tgz_file = tarfile.open(file_path) tgz_file.extractall(path=dir_path.parent) # the aclImdb.tgz has aclImdb dir packed assert dir_path.exists() print(f"Extracted to {dir_path}") CLAS_PATH = dir_path.parent CLAS_PATH.mkdir(exist_ok=True) LM_PATH = dir_path.parent /'imdb_lm' LM_PATH.mkdir(exist_ok=True) # processing the split files to create train.csv and test.csv in fastai format col_names = ['labels', 'text'] trn_texts, trn_labels = get_texts(dir_path/ 'train') val_texts, val_labels = get_texts(dir_path / 'test') np.random.seed(42) trn_idx = np.random.permutation(len(trn_texts)) val_idx = np.random.permutation(len(val_texts)) trn_texts = trn_texts[trn_idx] val_texts = val_texts[val_idx] trn_labels = trn_labels[trn_idx] val_labels = val_labels[val_idx] df_trn = pd.DataFrame({'text': trn_texts, 'labels': trn_labels}, columns=col_names) df_val = pd.DataFrame({'text': val_texts, 'labels': val_labels}, columns=col_names) print(f"df_trn has {len(df_trn)} rows, while df_val has {len(df_val)} rows") print(f"Writing them to {CLAS_PATH}") df_trn[df_trn['labels'] != 2].to_csv(CLAS_PATH / 'train.csv', header=False, index=False) df_val.to_csv(CLAS_PATH / 'test.csv', header=False, index=False) df_trn[df_trn['labels'] == 2].to_csv(CLAS_PATH / 'unsup.csv', header=False, index=False) (CLAS_PATH / 'classes.txt').open('w', encoding='utf-8').writelines(f'{o}\n' for o in CLASSES) 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] if split == TRN: file_path = file_path % lang elif lang == EN: 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: tokenizer = SentencePieceTokenizer(spm_path, use_moses=False, lang=lang) tok = tokenizer.tok_fun_with_sp(lang) tokenize = lambda x: tokenizer.process_text(x, tok) print("WARNING: Sentence Piece is not tested on XNLI yet") else: tokenize = lambda x: x.split(' ') toks, lbls = [], [] print(f'Reading {file_path}...') with open(file_path, encoding='utf-8') as f: reader = csv.reader(f, delimiter='\t', quoting=csv.QUOTE_NONE) for i, row in enumerate(reader): if i == 0: # skip the header continue # the examples are already tokenized with Moses if split == TRN: premise, hypo, label = row else: ex_lang = row[0] if ex_lang != lang: continue premise, hypo, label = row[-3], row[-2], row[1] # TODO add BOS premise_toks = tokenize(premise) + [EOS] hypo_toks = tokenize(hypo) + [EOS] toks.append(premise_toks + [SEP] + hypo_toks) lbls.append(label) return toks, lbls def read_clas_data(dir_path, dataset, lang) -> Tuple[Dict[str, List[List[str]]], Dict[str, List[str]]]: """ Read the dataset from the classification datasets and tokenize them. :param dir_path: the path to the dataset :param dataset: the name of the dataset :param lang: the language :return: a tuple consisting of: 1. a dictionary mapping splits to a list of lists of tokens 2. a dictionary mapping splits to a list of labels """ processors = { 'xnli': read_xnli } processor = processors[dataset] toks, lbls = {}, {} toks[TRN], lbls[TRN] = processor(dir_path, lang, TRN) toks[TST], lbls[TST] = processor(dir_path, lang, TST) if dataset == IMDB: # for IMDb, we need to split off a separate validation set # note that we train and fine-tune ULMFiT on the full training set in the paper # to do this, we can just keep the training set the same val_len = max(int(len(toks[TRN]) * 0.1), 2) # fastai does not work with validation set of size 1 trn_len = len(toks[TRN]) - val_len toks[TRN], toks[VAL] = toks[TRN][:trn_len], toks[TRN][trn_len:] lbls[TRN], lbls[VAL] = lbls[TRN][:trn_len], lbls[TRN][trn_len:] else: toks[VAL], lbls[VAL] = processor(dir_path, lang, VAL) return toks, lbls def replace_number(token): """Replaces a number and returns a list of one or multiple tokens.""" if number_match_re.match(token): return number_split_re.sub(r' @\1@ ', token) return token def read_file(file_path, outname=None): """Reads a text file and writes it to a .csv.""" with open(file_path, encoding='utf8') as f: text = f.readlines() df = pd.DataFrame( {'text': np.array(text), 'labels': np.zeros(len(text))}, columns=['labels', 'text']) if outname is not None: df.to_csv(file_path.parent / f'{outname}.csv', header=False, index=False) return df def read_whitespace_file(filepath): """Reads a file and prepares the tokens.""" tokens = [] with open(filepath, encoding='utf-8') as f: for line in f: # newlines are replaced with EOS tokens.append(line.split() + [EOS]) return np.array(tokens) class DataStump: """Placeholder class as LanguageModelLoader requires object with ids attribute.""" def __init__(self, ids): self.ids = ids self.loss_func = F.cross_entropy def validate(model, ids, bptt=2000): """ Return the validation loss and perplexity of a model :param model: model to test :param ids: data on which to evaluate the model :param bptt: bptt for this evaluation (doesn't change the result, only the speed) From https://github.com/sgugger/Adam-experiments/blob/master/lm_val_fns.py#L34 """ data = TextReader(np.concatenate(ids), bptt) model.eval() model.reset() total_loss, num_examples = 0., 0 for inputs, targets in tqdm(data): outputs, raws, outs = model(to_device(inputs, None)) p_vocab = F.softmax(outputs, 1) for i, pv in enumerate(p_vocab): targ_pred = pv[targets[i]] total_loss -= torch.log(targ_pred.detach()) num_examples += len(inputs) mean = total_loss / num_examples # divide by total number of tokens return mean, np.exp(mean) class TextReader(): """ Returns a language model iterator that iterates through batches that are of length N(bptt,5) The first batch returned is always bptt+25; the max possible width. This is done because of they way that pytorch allocates cuda memory in order to prevent multiple buffers from being created as the batch width grows. From: https://github.com/sgugger/Adam-experiments/blob/master/lm_val_fns.py#L3 """ def __init__(self, nums, bptt, backwards=False): self.bptt,self.backwards = bptt,backwards self.data = self.batchify(nums) self.i,self.iter = 0,0 self.n = len(self.data) def __iter__(self): self.i,self.iter = 0,0 while self.i < self.n-1 and self.iter