""" The enwiki/conll Dataset reader/provider using torchtext. The datasets were crated using the scripts from: https://github.com/samuelbroscheit/entity_knowledge_in_bert/tree/master/bert_entity/preprocessing The get_dataset.collate_batch function is influenced by: https://raw.githubusercontent.com/samuelbroscheit/entity_knowledge_in_bert/master/bert_entity/data_loader_wiki.py Please note that the pre-processed fine-tuning data will be automatically downloaded upon instantiation of the data readers and the result will be saved under /home//.cache/torch/text/datasets/ (in linux systems) The expected sizes of the auto-downloaded datasets: - Step 1 (general knowledge fine-tuning): enwiki-2023-spel-roberta-tokenized-aug-27-2023.tar.gz: 19.1 GBs * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * You can delete the file above once fine-tuning step 1 is done, and you are moving on to step 2. * * in the cleaning up process, make sure you remove the cached validation set files under .checkpoints * * directory as well * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * - Step 2 (general knowledge fine-tuning): enwiki-2023-spel-roberta-tokenized-aug-27-2023-retokenized.tar.gz: 17.5 GBs * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * You can delete the file above once fine-tuning step 2 is done, and you are moving on to step 3. * * in the cleaning up process, make sure you remove the cached validation set files under .checkpoints * * directory as well * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * - Step 3 (domain specific fine-tuning): aida-conll-spel-roberta-tokenized-aug-23-2023.tar.gz: 5.1 MBs No extra preprocessing step will be required, as soon as you start the fine-tuning script for each step, the proper fine-tuning dataset will be downloaded and will be served **without** the need for unzipping. """ import os import json import numpy from functools import partial from collections import OrderedDict from tqdm import tqdm from typing import Union, Tuple import torch from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler from torchdata.datapipes.iter import FileOpener, IterableWrapper, HttpReader, FileLister from torchtext.data.datasets_utils import _wrap_split_argument, _create_dataset_directory from torchtext.utils import download_from_url from transformers import AutoTokenizer, BatchEncoding from configuration import (get_aida_plus_wikipedia_plus_out_of_domain_vocab, get_aida_train_canonical_redirects, get_aida_vocab, get_ood_vocab, get_checkpoints_dir, get_base_model_name, device) BERT_MODEL_NAME = get_base_model_name() MAX_SPAN_ANNOTATION_SIZE = 4 class StaticAccess: def __init__(self): self.mentions_vocab, self.mentions_itos = None, None self.set_vocab_and_itos_to_all() self.aida_canonical_redirects = get_aida_train_canonical_redirects() self._all_vocab_mask_for_aida = None self._all_vocab_mask_for_ood = None def set_vocab_and_itos_to_all(self): self.mentions_vocab = get_aida_plus_wikipedia_plus_out_of_domain_vocab() self.mentions_itos = [w[0] for w in sorted(self.mentions_vocab.items(), key=lambda x: x[1])] @staticmethod def get_aida_vocab_and_itos(): aida_mentions_vocab = get_aida_vocab() aida_mentions_itos = [w[0] for w in sorted(aida_mentions_vocab.items(), key=lambda x: x[1])] return aida_mentions_vocab, aida_mentions_itos def shrink_vocab_to_aida(self): self.mentions_vocab, self.mentions_itos = self.get_aida_vocab_and_itos() def get_all_vocab_mask_for_aida(self): if self._all_vocab_mask_for_aida is None: mentions_vocab = get_aida_plus_wikipedia_plus_out_of_domain_vocab() mask = torch.ones(len(mentions_vocab)).to(device) mask = mask * -10000 mask[torch.Tensor([mentions_vocab[x] for x in get_aida_vocab()]).long()] = 0 self._all_vocab_mask_for_aida = mask return self._all_vocab_mask_for_aida def get_all_vocab_mask_for_ood(self): if self._all_vocab_mask_for_ood is None: mentions_vocab = get_aida_plus_wikipedia_plus_out_of_domain_vocab() mask = torch.ones(len(mentions_vocab)).to(device) mask = mask * -10000 mask[torch.Tensor([mentions_vocab[x] for x in get_ood_vocab()]).long()] = 0 self._all_vocab_mask_for_ood = mask return self._all_vocab_mask_for_ood dl_sa = StaticAccess() class ENWIKI20230827Config: URL = "https://1sfu-my.sharepoint.com/:u:/g/personal/sshavara_sfu_ca/Ea3IVbOpkTJKpASNyL9aFGMBQpH0ABU2hQa-wYyakkZ9TQ?e=DJFF3v&download=1" MD5 = "eb9a54a8f1f858cdcbf6c750942a896f" PATH = "enwiki-2023-spel-roberta-tokenized-aug-27-2023.tar.gz" DATASET_NAME = "WIKIPEDIA20230827" NUM_LINES = {'train': 3055221, 'valid': 1000, 'test': 1000} class ENWIKI20230827V2Config: URL = 'https://1sfu-my.sharepoint.com/:u:/g/personal/sshavara_sfu_ca/EeS_Tgl_CFJNiTh6YH5IDrsBocEZUsZV3lxPB6pleTxyxw?e=caH1cf&download=1' MD5 = "83a37f528800a463cd1a376e80ffc744" PATH = "enwiki-2023-spel-roberta-tokenized-aug-27-2023-retokenized.tar.gz" DATASET_NAME = "WIKIPEDIA20230827V2" NUM_LINES = {'train': 3038581, 'valid': 996} class AIDA20230827Config: URL = "https://1sfu-my.sharepoint.com/:u:/g/personal/sshavara_sfu_ca/EajEGYyf8LBOoxqDaiPBvbgBwFuEC08nssvZwGJWsG_HXg?e=wAwV6H&download=1" MD5 = "8078529d5df96d0d1ecf6a505fdb767a" PATH = "aida-conll-spel-roberta-tokenized-aug-23-2023.tar.gz" DATASET_NAME = "AIDA20230827" NUM_LINES = {'train': 1585, 'valid': 391, 'test': 372} tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME, cache_dir=get_checkpoints_dir() / "hf") WIKI_EXTRACTED_FILES = {"train": "train.json", "valid": "valid.json", "test": " test.json"} def wiki_filter_fn(split, fname_and_stream): return WIKI_EXTRACTED_FILES[split] in fname_and_stream[0] def wiki_data_record_convert(line): element = json.loads(line) r = {'tokens': element['tokens'], 'mentions': [], 'mention_entity_probs': [], 'mention_probs': []} for token, mentions, mention_entity_probs, mention_probs in zip(element['tokens'], element['mentions'], element['mention_entity_probs'], element['mention_probs']): if len(mention_probs) < len(mentions): mention_probs.extend([1.0 for _ in range(len(mentions) - len(mention_probs))]) sorted_mentions = sorted(list(zip(mentions, mention_entity_probs, mention_probs)), key=lambda x: x[1], reverse=True) mentions_ = [dl_sa.aida_canonical_redirects[x[0]] if x[0] in dl_sa.aida_canonical_redirects else x[0] for x in sorted_mentions if x[0]] # ignore mentions that are None mention_entity_probs_ = [x[1] for x in sorted_mentions if x[0]] # ignore prob. for None mentions mention_probs_ = [x[2] for x in sorted_mentions if x[0]] # ignore m_probs for None mentions r['mentions'].append(mentions_[:MAX_SPAN_ANNOTATION_SIZE]) r['mention_probs'].append(mention_probs_[:MAX_SPAN_ANNOTATION_SIZE]) r['mention_entity_probs'].append(mention_entity_probs_[:MAX_SPAN_ANNOTATION_SIZE]) if len(mentions_) > MAX_SPAN_ANNOTATION_SIZE: r['mention_entity_probs'][-1] = [x / sum(r['mention_entity_probs'][-1]) for x in r['mention_entity_probs'][-1]] return r @_create_dataset_directory(dataset_name=ENWIKI20230827Config.DATASET_NAME) @_wrap_split_argument(("train", "valid", "test")) def ENWIKI20230827(root: str, split: Union[Tuple[str], str]): root = root if root else ".data" path = root + "/" + ENWIKI20230827Config.PATH if not os.path.exists(path): download_from_url(ENWIKI20230827Config.URL, root=root, path=path, hash_value=ENWIKI20230827Config.MD5, hash_type='md5') online_reader_dp = FileLister(root, ENWIKI20230827Config.PATH) tar_file_dp = FileOpener(online_reader_dp, mode="b").load_from_tar().filter( partial(wiki_filter_fn, split)).readlines(return_path=False).map(wiki_data_record_convert) return tar_file_dp @_create_dataset_directory(dataset_name=ENWIKI20230827V2Config.DATASET_NAME) @_wrap_split_argument(("train", "valid")) def ENWIKI20230827V2(root: str, split: Union[Tuple[str], str]): root = root if root else ".data" path = root + "/" + ENWIKI20230827V2Config.PATH if not os.path.exists(path): download_from_url(ENWIKI20230827V2Config.URL, root=root, path=path, hash_value=ENWIKI20230827V2Config.MD5, hash_type='md5') online_reader_dp = FileLister(root, ENWIKI20230827V2Config.PATH) tar_file_dp = FileOpener(online_reader_dp, mode="b").load_from_tar().filter( partial(wiki_filter_fn, split)).readlines(return_path=False).map(wiki_data_record_convert) return tar_file_dp def aida_path_fn(r, _=None): return os.path.join(r, AIDA20230827Config.PATH) def aida_select_split(s, file_name_data): return file_name_data[1][s] def aida_data_record_convert(r): for x in r: # making sure each token comes with exactly one annotation assert len(x) == 7 or len(x) == 8 # whether it contains the candidates or not return {"tokens": [x[0] for x in r], "mentions": [[x[4] if x[4] else "|||O|||"] for x in r], "mention_entity_probs": [[1.0] for _ in r], "mention_probs": [[1.0] for _ in r], "candidates": [x[7] if x[7] else [] for x in r] if len(x) == 8 else [[] for x in r]} @_create_dataset_directory(dataset_name=AIDA20230827Config.DATASET_NAME) @_wrap_split_argument(('train', 'valid', 'test')) def AIDA20230827(root, split): online_reader_dp = HttpReader(IterableWrapper([AIDA20230827Config.URL])).on_disk_cache( filepath_fn=partial(aida_path_fn, root), hash_dict={aida_path_fn(root): AIDA20230827Config.MD5}, hash_type="md5").end_caching(mode="wb", same_filepath_fn=True) return FileOpener(online_reader_dp, mode="b").load_from_tar().parse_json_files().flatmap( partial(aida_select_split, split)).map(aida_data_record_convert) class DistributableDataset(Dataset): """ Based on the documentations in torch.utils.data.DataLoader, `IterableDataset` does not support custom `sampler` Therefore we cannot use the DistributedSampler with the DataLoader to split the data samples. This class is a workaround to make the IterableDataset work with the DistributedSampler. """ def __init__(self, dataset, size, world_size, rank): self.size = size self.data = iter(dataset) self.world_size = world_size self.rank = rank self.initial_fetch = True def __len__(self): return self.size def __getitem__(self, index): # Since we don't have a means of accessing the data by indices, we try skipping the indices that we believe # belong to other processes skip_size = self.rank if self.initial_fetch else self.world_size - 1 self.initial_fetch = False for _ in range(skip_size): next(self.data) return next(self.data) def convert_is_in_mention_to_bioes(is_in_mention): # B = 0, I = 1, O = 2, E = 3, S = 4 bioes = [] for iim, current in enumerate(is_in_mention): before = is_in_mention[iim - 1] if iim > 0 else 0 after = is_in_mention[iim + 1] if iim < len(is_in_mention) - 1 else 0 bioes.append( 2 if not current else (4 if not before and not after else (0 if not before else (3 if not after else 1)))) return bioes def get_dataset(dataset_name: str, split: str, batch_size: int, get_labels_with_high_model_score=None, label_size: int = 0, load_distributed: bool = False, world_size: int = 1, rank: int = 0, use_retokenized_wikipedia_data: bool = False): """ :param dataset_name: The dataset name can either be "enwiki" or "aida" :param split: the requested dataset split which can be 'train', 'valid' or 'test' :param batch_size: the size of the resulting batch from the data loader :param get_labels_with_high_model_score: The function that finds high scoring negative samples for the model :param label_size: The maximum output distribution size. You can pass the output vocabulary size for this parameter. :param load_distributed: The flag hinting whether the data loader will be loaded in a multi-gpu setting. :param world_size: the number of machines that the dataloader is expected to serve. :param rank: the rank of the gpu on which the data is expected to be served. :param use_retokenized_wikipedia_data: a flag indicating whether to use ENWIKI20230827 dataset or ENWIKI20230827V2 """ assert dataset_name in ["enwiki", "aida"] if not load_distributed or rank == 0: print(f"Initializing the {dataset_name.upper()}/{split} dataset ...") def collate_batch(batch): data = {} for key in ["tokens", "mentions", "mention_entity_probs", "eval_mask", "candidates", "is_in_mention", "bioes"]: data[key] = [] for annotated_line_in_file in batch: data["tokens"].append(tokenizer.convert_tokens_to_ids(annotated_line_in_file["tokens"])) data["mentions"].append([ [(dl_sa.mentions_vocab[x] if x not in dl_sa.aida_canonical_redirects else dl_sa.mentions_vocab[dl_sa.aida_canonical_redirects[x]]) if x is not None and x not in ['Gmina_Żabno'] else dl_sa.mentions_vocab["|||O|||"] for x in el] for el in annotated_line_in_file["mentions"] ]) data["mention_entity_probs"].append(annotated_line_in_file["mention_entity_probs"]) data["eval_mask"].append(list(map( lambda item: 1 if len(item) == 1 else 0, annotated_line_in_file["mention_probs"]))) is_in_mention = [1 if x != '|||O|||' else 0 for el, elp in zip( annotated_line_in_file["mentions"], annotated_line_in_file["mention_entity_probs"]) for x, y in zip(el, elp) if y == max(elp)] data["is_in_mention"].append(is_in_mention) data["bioes"].append(convert_is_in_mention_to_bioes(is_in_mention)) maxlen = max([len(x) for x in data["tokens"]]) token_ids = torch.LongTensor([sample + [0] * (maxlen - len(sample)) for sample in data["tokens"]]) eval_mask = torch.LongTensor([sample + [0] * (maxlen - len(sample)) for sample in data["eval_mask"]]) is_in_mention = torch.LongTensor([sample + [0] * (maxlen - len(sample)) for sample in data["is_in_mention"]]) bioes = torch.LongTensor([sample + [2] * (maxlen - len(sample)) for sample in data["bioes"]]) if get_labels_with_high_model_score: labels_with_high_model_score = get_labels_with_high_model_score(token_ids) else: labels_with_high_model_score = None subword_mentions = create_output_with_negative_examples( data["mentions"], data["mention_entity_probs"], token_ids.size(0), token_ids.size(1), len(dl_sa.mentions_vocab), label_size, labels_with_high_model_score) inputs = BatchEncoding({ 'token_ids': token_ids, 'eval_mask': eval_mask, 'raw_mentions': data["mentions"], 'is_in_mention': is_in_mention, "bioes": bioes }) return inputs, subword_mentions if not load_distributed or rank == 0: print(f"Done initializing the {dataset_name.upper()}/{split} dataset ...") wikipedia_dataset = ENWIKI20230827 wikipedia_dataset_config = ENWIKI20230827Config retokenized_wikipedia_dataset = ENWIKI20230827V2 retokenized_wikipedia_dataset_config = ENWIKI20230827V2Config aida_dataset = AIDA20230827 aida_dataset_config = AIDA20230827Config dset_class = (retokenized_wikipedia_dataset if use_retokenized_wikipedia_data else wikipedia_dataset) \ if dataset_name == "enwiki" else aida_dataset d_size = (retokenized_wikipedia_dataset_config.NUM_LINES[split] if use_retokenized_wikipedia_data else wikipedia_dataset_config.NUM_LINES[split]) \ if dataset_name == "enwiki" else aida_dataset_config.NUM_LINES[split] dataset_ = DistributableDataset(dset_class(split=split, root=get_checkpoints_dir()), d_size, world_size, rank) \ if load_distributed else dset_class(split=split, root=get_checkpoints_dir()) return DataLoader(dataset_, batch_size=batch_size, collate_fn=collate_batch, sampler=DistributedSampler(dataset_, num_replicas=world_size, rank=rank)) \ if load_distributed and split == "train" else DataLoader(dset_class(split=split, root=get_checkpoints_dir()), batch_size=batch_size, collate_fn=collate_batch) def create_output_with_negative_examples(batch_entity_ids, batch_entity_probs, batch_size, maxlen, label_vocab_size, label_size, labels_with_high_model_score=None): all_entity_ids = OrderedDict() for batch_offset, (batch_item_token_item_entity_ids, batch_item_token_entity_probs) in enumerate( zip(batch_entity_ids, batch_entity_probs) ): for tok_id, (token_entity_ids, token_entity_probs) in enumerate( zip(batch_item_token_item_entity_ids, batch_item_token_entity_probs) ): for eid in token_entity_ids: if eid not in all_entity_ids: all_entity_ids[eid] = len(all_entity_ids) # ##################################################### shared_label_ids = list(all_entity_ids.keys()) if len(shared_label_ids) < label_size and labels_with_high_model_score is not None: negative_examples = set(labels_with_high_model_score) negative_examples.difference_update(shared_label_ids) shared_label_ids += list(negative_examples) if len(shared_label_ids) < label_size: negative_samples = set(numpy.random.choice(label_vocab_size, label_size, replace=False)) negative_samples.difference_update(shared_label_ids) shared_label_ids += list(negative_samples) shared_label_ids = shared_label_ids[: label_size] all_batch_entity_ids, batch_shared_label_ids = all_entity_ids, shared_label_ids if label_size > 0: label_probs = torch.zeros(batch_size, maxlen, len(batch_shared_label_ids)) else: label_probs = torch.zeros(batch_size, maxlen, label_vocab_size) # loop through the batch x tokens x (label_ids, label_probs) for batch_offset, (batch_item_token_item_entity_ids, batch_item_token_entity_probs) in enumerate( zip(batch_entity_ids, batch_entity_probs) ): # loop through tokens x (label_ids, label_probs) for tok_id, (token_entity_ids, token_entity_probs) in enumerate( zip(batch_item_token_item_entity_ids, batch_item_token_entity_probs)): if label_size is None: label_probs[batch_offset][tok_id][torch.LongTensor(token_entity_ids)] = torch.Tensor( batch_item_token_item_entity_ids) else: label_probs[batch_offset][tok_id][ torch.LongTensor(list(map(all_batch_entity_ids.__getitem__, token_entity_ids))) ] = torch.Tensor(token_entity_probs) label_ids = torch.LongTensor(batch_shared_label_ids) return BatchEncoding({ "ids": label_ids, # of size label_size "probs": label_probs, # of size input_batch_size x input_max_len x label_size "dictionary": {v: k for k, v in all_batch_entity_ids.items()} # contains all original ids for mentions in batch }) def _make_vocab_file(): wiki_vocab = set() vocab_file = open("enwiki_20230827.txt", "w") for spl in ['train', 'valid', 'test']: for el in tqdm(ENWIKI20230827(split=spl, root=get_checkpoints_dir())): for x in el['mentions']: for y in x: if y not in wiki_vocab: vocab_file.write(f"{y}\n") wiki_vocab.add(y) vocab_file.close()