import logging from typing import ( Any, Callable, Dict, Generator, Iterator, List, NamedTuple, Optional, Tuple, Union, ) import numpy as np import torch from reader.data.relik_reader_data_utils import ( add_noise_to_value, batchify, batchify_matrices, batchify_tensor, chunks, flatten, ) from reader.data.relik_reader_sample import RelikReaderSample, load_relik_reader_samples from torch.utils.data import IterableDataset from transformers import AutoTokenizer from relik.reader.utils.special_symbols import NME_SYMBOL logger = logging.getLogger(__name__) class TokenizationOutput(NamedTuple): input_ids: torch.Tensor attention_mask: torch.Tensor token_type_ids: torch.Tensor prediction_mask: torch.Tensor special_symbols_mask: torch.Tensor special_symbols_mask_entities: torch.Tensor class RelikREDataset(IterableDataset): def __init__( self, dataset_path: str, materialize_samples: bool, transformer_model: str, special_symbols: List[str], shuffle_candidates: Optional[Union[bool, float]], flip_candidates: Optional[Union[bool, float]], relations_definitions: Union[str, Dict[str, str]], for_inference: bool, entities_definitions: Optional[Union[str, Dict[str, str]]] = None, special_symbols_entities: Optional[List[str]] = None, noise_param: float = 0.1, sorting_fields: Optional[str] = None, tokens_per_batch: int = 2048, batch_size: int = None, max_batch_size: int = 128, section_size: int = 50_000, prebatch: bool = True, max_candidates: int = 0, add_gold_candidates: bool = True, use_nme: bool = True, min_length: int = 5, max_length: int = 2048, model_max_length: int = 1000, skip_empty_training_samples: bool = True, drop_last: bool = False, samples: Optional[Iterator[RelikReaderSample]] = None, **kwargs, ): super().__init__(**kwargs) self.dataset_path = dataset_path self.materialize_samples = materialize_samples self.samples: Optional[List[RelikReaderSample]] = None if self.materialize_samples: self.samples = list() self.tokenizer = self._build_tokenizer(transformer_model, special_symbols) self.special_symbols = special_symbols self.special_symbols_entities = special_symbols_entities self.shuffle_candidates = shuffle_candidates self.flip_candidates = flip_candidates self.for_inference = for_inference self.noise_param = noise_param self.batching_fields = ["input_ids"] self.sorting_fields = ( sorting_fields if sorting_fields is not None else self.batching_fields ) # open relations definitions file if needed if type(relations_definitions) == str: relations_definitions = { line.split("\t")[0]: line.split("\t")[1] for line in open(relations_definitions) } self.max_candidates = max_candidates self.relations_definitions = relations_definitions self.entities_definitions = entities_definitions self.add_gold_candidates = add_gold_candidates self.use_nme = use_nme self.min_length = min_length self.max_length = max_length self.model_max_length = ( model_max_length if model_max_length < self.tokenizer.model_max_length else self.tokenizer.model_max_length ) self.transformer_model = transformer_model self.skip_empty_training_samples = skip_empty_training_samples self.drop_last = drop_last self.samples = samples self.tokens_per_batch = tokens_per_batch self.batch_size = batch_size self.max_batch_size = max_batch_size self.section_size = section_size self.prebatch = prebatch def _build_tokenizer(self, transformer_model: str, special_symbols: List[str]): return AutoTokenizer.from_pretrained( transformer_model, additional_special_tokens=[ss for ss in special_symbols], add_prefix_space=True, ) @property def fields_batcher(self) -> Dict[str, Union[None, Callable[[list], Any]]]: fields_batchers = { "input_ids": lambda x: batchify( x, padding_value=self.tokenizer.pad_token_id ), "attention_mask": lambda x: batchify(x, padding_value=0), "token_type_ids": lambda x: batchify(x, padding_value=0), "prediction_mask": lambda x: batchify(x, padding_value=1), "global_attention": lambda x: batchify(x, padding_value=0), "token2word": None, "sample": None, "special_symbols_mask": lambda x: batchify(x, padding_value=False), "special_symbols_mask_entities": lambda x: batchify(x, padding_value=False), "start_labels": lambda x: batchify(x, padding_value=-100), "end_labels": lambda x: batchify_matrices(x, padding_value=-100), "disambiguation_labels": lambda x: batchify(x, padding_value=-100), "relation_labels": lambda x: batchify_tensor(x, padding_value=-100), "predictable_candidates": None, } if "roberta" in self.transformer_model: del fields_batchers["token_type_ids"] return fields_batchers def _build_input_ids( self, sentence_input_ids: List[int], candidates_input_ids: List[List[int]] ) -> List[int]: return ( [self.tokenizer.cls_token_id] + sentence_input_ids + [self.tokenizer.sep_token_id] + flatten(candidates_input_ids) + [self.tokenizer.sep_token_id] ) def _get_special_symbols_mask(self, input_ids: torch.Tensor) -> torch.Tensor: special_symbols_mask = input_ids >= ( len(self.tokenizer) - len(self.special_symbols + self.special_symbols_entities) ) special_symbols_mask[0] = True return special_symbols_mask def _build_tokenizer_essentials( self, input_ids, original_sequence ) -> TokenizationOutput: input_ids = torch.tensor(input_ids, dtype=torch.long) attention_mask = torch.ones_like(input_ids) total_sequence_len = len(input_ids) predictable_sentence_len = len(original_sequence) # token type ids token_type_ids = torch.cat( [ input_ids.new_zeros( predictable_sentence_len + 2 ), # original sentence bpes + CLS and SEP input_ids.new_ones(total_sequence_len - predictable_sentence_len - 2), ] ) # prediction mask -> boolean on tokens that are predictable prediction_mask = torch.tensor( [1] + ([0] * predictable_sentence_len) + ([1] * (total_sequence_len - predictable_sentence_len - 1)) ) assert len(prediction_mask) == len(input_ids) # special symbols mask special_symbols_mask = input_ids >= ( len(self.tokenizer) - len(self.special_symbols) # + self.special_symbols_entities) ) if self.entities_definitions is not None: # select only the first N true values where N is len(entities_definitions) special_symbols_mask_entities = special_symbols_mask.clone() special_symbols_mask_entities[ special_symbols_mask_entities.cumsum(0) > len(self.entities_definitions) ] = False special_symbols_mask = special_symbols_mask ^ special_symbols_mask_entities else: special_symbols_mask_entities = special_symbols_mask.clone() return TokenizationOutput( input_ids, attention_mask, token_type_ids, prediction_mask, special_symbols_mask, special_symbols_mask_entities, ) def _build_labels( self, sample, tokenization_output: TokenizationOutput, ) -> Tuple[torch.Tensor, torch.Tensor]: start_labels = [0] * len(tokenization_output.input_ids) end_labels = [] sample.entities.sort(key=lambda x: (x[0], x[1])) prev_start_bpe = -1 num_repeat_start = 0 if self.entities_definitions: sample.entities = [(ce[0], ce[1], ce[2]) for ce in sample.entities] sample.entity_candidates = list(self.entities_definitions.keys()) disambiguation_labels = torch.zeros( len(sample.entities), len(sample.entity_candidates) + len(sample.candidates), ) else: sample.entities = [(ce[0], ce[1], "") for ce in sample.entities] disambiguation_labels = torch.zeros( len(sample.entities), len(sample.candidates) ) ignored_labels_indices = tokenization_output.prediction_mask == 1 for idx, c_ent in enumerate(sample.entities): start_bpe = sample.word2token[c_ent[0]][0] + 1 end_bpe = sample.word2token[c_ent[1] - 1][-1] + 1 class_index = idx start_labels[start_bpe] = class_index + 1 # +1 for the NONE class if start_bpe != prev_start_bpe: end_labels.append([0] * len(tokenization_output.input_ids)) # end_labels[-1][:start_bpe] = [-100] * start_bpe end_labels[-1][end_bpe] = class_index + 1 else: end_labels[-1][end_bpe] = class_index + 1 num_repeat_start += 1 if self.entities_definitions: entity_type_idx = sample.entity_candidates.index(c_ent[2]) disambiguation_labels[idx, entity_type_idx] = 1 prev_start_bpe = start_bpe start_labels = torch.tensor(start_labels, dtype=torch.long) start_labels[ignored_labels_indices] = -100 end_labels = torch.tensor(end_labels, dtype=torch.long) end_labels[ignored_labels_indices.repeat(len(end_labels), 1)] = -100 relation_labels = torch.zeros( len(sample.entities), len(sample.entities), len(sample.candidates) ) # sample.relations = [] for re in sample.triplets: if re["relation"]["name"] not in sample.candidates: re_class_index = len(sample.candidates) - 1 else: re_class_index = sample.candidates.index( re["relation"]["name"] ) # should remove this +1 if self.entities_definitions: subject_class_index = sample.entities.index( ( re["subject"]["start"], re["subject"]["end"], re["subject"]["type"], ) ) object_class_index = sample.entities.index( (re["object"]["start"], re["object"]["end"], re["object"]["type"]) ) else: subject_class_index = sample.entities.index( (re["subject"]["start"], re["subject"]["end"], "") ) object_class_index = sample.entities.index( (re["object"]["start"], re["object"]["end"], "") ) relation_labels[subject_class_index, object_class_index, re_class_index] = 1 if self.entities_definitions: disambiguation_labels[ subject_class_index, re_class_index + len(sample.entity_candidates) ] = 1 disambiguation_labels[ object_class_index, re_class_index + len(sample.entity_candidates) ] = 1 # sample.relations.append([re['subject']['start'], re['subject']['end'], re['subject']['type'], re['relation']['name'], re['object']['start'], re['object']['end'], re['object']['type']]) else: disambiguation_labels[subject_class_index, re_class_index] = 1 disambiguation_labels[object_class_index, re_class_index] = 1 # sample.relations.append([re['subject']['start'], re['subject']['end'], "", re['relation']['name'], re['object']['start'], re['object']['end'], ""]) return start_labels, end_labels, disambiguation_labels, relation_labels def __iter__(self): dataset_iterator = self.dataset_iterator_func() current_dataset_elements = [] i = None for i, dataset_elem in enumerate(dataset_iterator, start=1): if ( self.section_size is not None and len(current_dataset_elements) == self.section_size ): for batch in self.materialize_batches(current_dataset_elements): yield batch current_dataset_elements = [] current_dataset_elements.append(dataset_elem) if i % 50_000 == 0: logger.info(f"Processed: {i} number of elements") if len(current_dataset_elements) != 0: for batch in self.materialize_batches(current_dataset_elements): yield batch if i is not None: logger.info(f"Dataset finished: {i} number of elements processed") else: logger.warning("Dataset empty") def dataset_iterator_func(self): data_samples = ( load_relik_reader_samples(self.dataset_path) if self.samples is None else self.samples ) for sample in data_samples: # input sentence tokenization input_tokenized = self.tokenizer( sample.tokens, return_offsets_mapping=True, add_special_tokens=False, is_split_into_words=True, ) input_subwords = input_tokenized["input_ids"] offsets = input_tokenized["offset_mapping"] token2word = [] word2token = {} count = 0 for i, offset in enumerate(offsets): if offset[0] == 0: token2word.append(i - count) word2token[i - count] = [i] else: token2word.append(token2word[-1]) word2token[token2word[-1]].append(i) count += 1 sample.token2word = token2word sample.word2token = word2token # input_subwords = sample.tokens[1:-1] # removing special tokens candidates_symbols = self.special_symbols if self.max_candidates > 0: # truncate candidates sample.candidates = sample.candidates[: self.max_candidates] # add NME as a possible candidate if self.use_nme: sample.candidates.insert(0, NME_SYMBOL) # training time sample mods if not self.for_inference: # check whether the sample has labels if not skip if ( sample.triplets is None or len(sample.triplets) == 0 ) and self.skip_empty_training_samples: logger.warning( "Sample {} has no labels, skipping".format(sample.sample_id) ) continue # add gold candidates if missing if self.add_gold_candidates: candidates_set = set(sample.candidates) candidates_to_add = [] for candidate_title in sample.triplets: if candidate_title["relation"]["name"] not in candidates_set: candidates_to_add.append( candidate_title["relation"]["name"] ) if len(candidates_to_add) > 0: # replacing last candidates with the gold ones # this is done in order to preserve the ordering added_gold_candidates = 0 gold_candidates_titles_set = set( set(ct["relation"]["name"] for ct in sample.triplets) ) for i in reversed(range(len(sample.candidates))): if ( sample.candidates[i] not in gold_candidates_titles_set and sample.candidates[i] != NME_SYMBOL ): sample.candidates[i] = candidates_to_add[ added_gold_candidates ] added_gold_candidates += 1 if len(candidates_to_add) == added_gold_candidates: break candidates_still_to_add = ( len(candidates_to_add) - added_gold_candidates ) while ( len(sample.candidates) <= len(candidates_symbols) and candidates_still_to_add != 0 ): sample.candidates.append( candidates_to_add[added_gold_candidates] ) added_gold_candidates += 1 candidates_still_to_add -= 1 # shuffle candidates if ( isinstance(self.shuffle_candidates, bool) and self.shuffle_candidates ) or ( isinstance(self.shuffle_candidates, float) and np.random.uniform() < self.shuffle_candidates ): np.random.shuffle(sample.candidates) if NME_SYMBOL in sample.candidates: sample.candidates.remove(NME_SYMBOL) sample.candidates.insert(0, NME_SYMBOL) # flip candidates if ( isinstance(self.flip_candidates, bool) and self.flip_candidates ) or ( isinstance(self.flip_candidates, float) and np.random.uniform() < self.flip_candidates ): for i in range(len(sample.candidates) - 1): if np.random.uniform() < 0.5: sample.candidates[i], sample.candidates[i + 1] = ( sample.candidates[i + 1], sample.candidates[i], ) if NME_SYMBOL in sample.candidates: sample.candidates.remove(NME_SYMBOL) sample.candidates.insert(0, NME_SYMBOL) # candidates encoding candidates_symbols = candidates_symbols[: len(sample.candidates)] relations_defs = [ "{} {}".format(cs, self.relations_definitions[ct]) if ct != NME_SYMBOL else NME_SYMBOL for cs, ct in zip(candidates_symbols, sample.candidates) ] if self.entities_definitions is not None: candidates_entities_symbols = list(self.special_symbols_entities) candidates_entities_symbols = candidates_entities_symbols[ : len(self.entities_definitions) ] entity_defs = [ "{} {}".format(cs, self.entities_definitions[ct]) for cs, ct in zip( candidates_entities_symbols, self.entities_definitions.keys() ) ] relations_defs = ( entity_defs + [self.tokenizer.sep_token] + relations_defs ) candidates_encoding_result = self.tokenizer.batch_encode_plus( relations_defs, add_special_tokens=False, ).input_ids # drop candidates if the number of input tokens is too long for the model if ( sum(map(len, candidates_encoding_result)) + len(input_subwords) + 20 # + 20 special tokens > self.model_max_length ): if self.for_inference: acceptable_tokens_from_candidates = ( self.model_max_length - 20 - len(input_subwords) ) while ( cum_len + len(candidates_encoding_result[i]) < acceptable_tokens_from_candidates ): cum_len += len(candidates_encoding_result[i]) i += 1 candidates_encoding_result = candidates_encoding_result[:i] if self.entities_definitions is not None: candidates_symbols = candidates_symbols[ : i - len(self.entities_definitions) ] sample.candidates = sample.candidates[ : i - len(self.entities_definitions) ] else: candidates_symbols = candidates_symbols[:i] sample.candidates = sample.candidates[:i] else: gold_candidates_set = set( [wl["relation"]["name"] for wl in sample.triplets] ) gold_candidates_indices = [ i for i, wc in enumerate(sample.candidates) if wc in gold_candidates_set ] if self.entities_definitions is not None: gold_candidates_indices = [ i + len(self.entities_definitions) for i in gold_candidates_indices ] # add entities indices gold_candidates_indices = gold_candidates_indices + list( range(len(self.entities_definitions)) ) necessary_taken_tokens = sum( map( len, [ candidates_encoding_result[i] for i in gold_candidates_indices ], ) ) acceptable_tokens_from_candidates = ( self.model_max_length - 20 - len(input_subwords) - necessary_taken_tokens ) assert acceptable_tokens_from_candidates > 0 i = 0 cum_len = 0 while ( cum_len + len(candidates_encoding_result[i]) < acceptable_tokens_from_candidates ): if i not in gold_candidates_indices: cum_len += len(candidates_encoding_result[i]) i += 1 new_indices = sorted( list(set(list(range(i)) + gold_candidates_indices)) ) np.random.shuffle(new_indices) candidates_encoding_result = [ candidates_encoding_result[i] for i in new_indices ] if self.entities_definitions is not None: sample.candidates = [ sample.candidates[i - len(self.entities_definitions)] for i in new_indices ] candidates_symbols = candidates_symbols[ : i - len(self.entities_definitions) ] else: candidates_symbols = [ candidates_symbols[i] for i in new_indices ] sample.window_candidates = [ sample.window_candidates[i] for i in new_indices ] if len(sample.candidates) == 0: logger.warning( "Sample {} has no candidates after truncation due to max length".format( sample.sample_id ) ) continue # final input_ids build input_ids = self._build_input_ids( sentence_input_ids=input_subwords, candidates_input_ids=candidates_encoding_result, ) # complete input building (e.g. attention / prediction mask) tokenization_output = self._build_tokenizer_essentials( input_ids, input_subwords ) # labels creation start_labels, end_labels, disambiguation_labels, relation_labels = ( None, None, None, None, ) if sample.entities is not None and len(sample.entities) > 0: ( start_labels, end_labels, disambiguation_labels, relation_labels, ) = self._build_labels( sample, tokenization_output, ) yield { "input_ids": tokenization_output.input_ids, "attention_mask": tokenization_output.attention_mask, "token_type_ids": tokenization_output.token_type_ids, "prediction_mask": tokenization_output.prediction_mask, "special_symbols_mask": tokenization_output.special_symbols_mask, "special_symbols_mask_entities": tokenization_output.special_symbols_mask_entities, "sample": sample, "start_labels": start_labels, "end_labels": end_labels, "disambiguation_labels": disambiguation_labels, "relation_labels": relation_labels, "predictable_candidates": candidates_symbols, } def preshuffle_elements(self, dataset_elements: List): # This shuffling is done so that when using the sorting function, # if it is deterministic given a collection and its order, we will # make the whole operation not deterministic anymore. # Basically, the aim is not to build every time the same batches. if not self.for_inference: dataset_elements = np.random.permutation(dataset_elements) sorting_fn = ( lambda elem: add_noise_to_value( sum(len(elem[k]) for k in self.sorting_fields), noise_param=self.noise_param, ) if not self.for_inference else sum(len(elem[k]) for k in self.sorting_fields) ) dataset_elements = sorted(dataset_elements, key=sorting_fn) if self.for_inference: return dataset_elements ds = list(chunks(dataset_elements, 64)) # todo: modified np.random.shuffle(ds) return flatten(ds) def materialize_batches( self, dataset_elements: List[Dict[str, Any]] ) -> Generator[Dict[str, Any], None, None]: if self.prebatch: dataset_elements = self.preshuffle_elements(dataset_elements) current_batch = [] # function that creates a batch from the 'current_batch' list def output_batch() -> Dict[str, Any]: assert ( len( set([len(elem["predictable_candidates"]) for elem in current_batch]) ) == 1 ), " ".join( map( str, [len(elem["predictable_candidates"]) for elem in current_batch] ) ) batch_dict = dict() de_values_by_field = { fn: [de[fn] for de in current_batch if fn in de] for fn in self.fields_batcher } # in case you provide fields batchers but in the batch # there are no elements for that field de_values_by_field = { fn: fvs for fn, fvs in de_values_by_field.items() if len(fvs) > 0 } assert len(set([len(v) for v in de_values_by_field.values()])) # todo: maybe we should report the user about possible # fields filtering due to "None" instances de_values_by_field = { fn: fvs for fn, fvs in de_values_by_field.items() if all([fv is not None for fv in fvs]) } for field_name, field_values in de_values_by_field.items(): field_batch = ( self.fields_batcher[field_name](field_values) if self.fields_batcher[field_name] is not None else field_values ) batch_dict[field_name] = field_batch return batch_dict max_len_discards, min_len_discards = 0, 0 should_token_batch = self.batch_size is None curr_pred_elements = -1 for de in dataset_elements: if ( should_token_batch and self.max_batch_size != -1 and len(current_batch) == self.max_batch_size ) or (not should_token_batch and len(current_batch) == self.batch_size): yield output_batch() current_batch = [] curr_pred_elements = -1 # todo support max length (and min length) as dicts too_long_fields = [ k for k in de if self.max_length != -1 and torch.is_tensor(de[k]) and len(de[k]) > self.max_length ] if len(too_long_fields) > 0: max_len_discards += 1 continue too_short_fields = [ k for k in de if self.min_length != -1 and torch.is_tensor(de[k]) and len(de[k]) < self.min_length ] if len(too_short_fields) > 0: min_len_discards += 1 continue if should_token_batch: de_len = sum(len(de[k]) for k in self.batching_fields) future_max_len = max( de_len, max( [ sum(len(bde[k]) for k in self.batching_fields) for bde in current_batch ], default=0, ), ) future_tokens_per_batch = future_max_len * (len(current_batch) + 1) num_predictable_candidates = len(de["predictable_candidates"]) if len(current_batch) > 0 and ( future_tokens_per_batch >= self.tokens_per_batch or ( num_predictable_candidates != curr_pred_elements and curr_pred_elements != -1 ) ): yield output_batch() current_batch = [] current_batch.append(de) curr_pred_elements = len(de["predictable_candidates"]) if len(current_batch) != 0 and not self.drop_last: yield output_batch() if max_len_discards > 0: if self.for_inference: logger.warning( f"WARNING: Inference mode is True but {max_len_discards} samples longer than max length were " f"found. The {max_len_discards} samples will be DISCARDED. If you are doing some kind of evaluation" f", this can INVALIDATE results. This might happen if the max length was not set to -1 or if the " f"sample length exceeds the maximum length supported by the current model." ) else: logger.warning( f"During iteration, {max_len_discards} elements were " f"discarded since longer than max length {self.max_length}" ) if min_len_discards > 0: if self.for_inference: logger.warning( f"WARNING: Inference mode is True but {min_len_discards} samples shorter than min length were " f"found. The {min_len_discards} samples will be DISCARDED. If you are doing some kind of evaluation" f", this can INVALIDATE results. This might happen if the min length was not set to -1 or if the " f"sample length is shorter than the minimum length supported by the current model." ) else: logger.warning( f"During iteration, {min_len_discards} elements were " f"discarded since shorter than min length {self.min_length}" ) def main(): special_symbols = [NME_SYMBOL] + [f"R-{i}" for i in range(50)] relik_dataset = RelikREDataset( "/home/huguetcabot/alby-re/alby/data/nyt-alby+/valid.jsonl", materialize_samples=False, transformer_model="microsoft/deberta-v3-base", special_symbols=special_symbols, shuffle_candidates=False, flip_candidates=False, for_inference=True, ) for batch in relik_dataset: print(batch) exit(0) if __name__ == "__main__": main()