def test_masked_class_logsoftmax_basic(self): # shape batch x M x K # model outputs preds = torch.tensor([[[2.0, 2.0, 1.0], [3.0, 5.0, 4.0]]]) # all that matters for this test is that the below is non-negative # since negative indicates masking entity_ids = torch.tensor([[[1, 3, 4], [5, 3, 1]]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax(pred=preds, mask=mask) torch_logsoftmax = torch.nn.LogSoftmax(dim=2) torch_log_preds = torch_logsoftmax(preds) assert torch.allclose(torch_log_preds, pred_log_preds) # if we mask one of the candidates, we should no longer # get the same result as torch fn which doesn't mask entity_ids = torch.tensor([[[1, 3, 4], [5, 3, -1]]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax(pred=preds, mask=mask) assert not torch.allclose(torch_log_preds, pred_log_preds) # make sure masked values are approximately zero before log (when exponented) assert torch.allclose( torch.tensor([[[0.422319, 0.422319, 0.155362], [0.119203, 0.880797, 0.0]]]), torch.exp(pred_log_preds), )
def disambig_loss(intermediate_output_dict, Y, active): """Returns the entity disambiguation loss on prediction heads. Args: intermediate_output_dict: output dict from the Emmental task flor Y: gold labels active: whether examples are "active" or not (used in Emmental slicing) Returns: loss """ # Grab the first value of training (when doing distributed training, we will have one per process) training = intermediate_output_dict[PRED_LAYER]["training"][0].item() assert type(training) is bool outs = intermediate_output_dict[PRED_LAYER]["final_scores"][DISAMBIG] mask = intermediate_output_dict["_input_"]["entity_cand_eid_mask"][active] labels = Y[active] # During eval, even if our model does not predict a NIC candidate, we allow for a NIC gold QID # This qid gets assigned the label of -2 and is always incorrect # As NLLLoss assumes classes of 0 to #classes-1 except for pad idx, we manually mask # the -2 labels for the loss computation only. As this is just for eval, it won't matter. masked_labels = labels if not training: label_mask = labels == -2 masked_labels = torch.where(~label_mask, labels, torch.ones_like(labels) * -1) temp = 0 for out in outs.values(): # batch x M x K -> transpose -> swap K classes with M spans for "k-dimensional" NLLloss log_probs = eval_utils.masked_class_logsoftmax(pred=out[active], mask=~mask).transpose( 1, 2) temp += nn.NLLLoss(ignore_index=-1)(log_probs, masked_labels.long()) return temp
def test_masked_class_logsoftmax_grads_excluded_alias(self): preds = torch.tensor([[[2.0, 4.0], [1.0, 4.0], [8.0, 2.0]]], requires_grad=True) # batch x M x K entity_ids = torch.tensor([[[1, -1], [4, 5], [8, 9]]]) # batch x M true_entity_class = torch.tensor([[0, -1, 1]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax( pred=preds, mask=mask).transpose(1, 2) pred_loss = torch.nn.NLLLoss(ignore_index=-1)(pred_log_preds, true_entity_class) pred_loss.backward() actual_grad = preds.grad true_entity_class_expanded = true_entity_class.unsqueeze(-1).expand_as( entity_ids) masked_actual_grad = torch.where( (entity_ids != -1) & (true_entity_class_expanded != -1), torch.ones_like(preds), actual_grad, ) # just put 1's where we want non-zeros and use mask above to only compare padded gradients expected_grad = torch.tensor([[[1.0, 0.0], [0.0, 0.0], [1.0, 1.0]]]) assert torch.allclose(expected_grad, masked_actual_grad)
def disambig_loss(self, training, outs, true_label, mask): """ Returns the entity disambiguation loss on prediction heads. """ device = next(iter(outs.values()))[0].device loss = LossPackage(device) for i, (loss_head, out) in enumerate(outs.items()): if FINAL_LOSS in loss_head: true_label_head = true_label[FINAL_LOSS] else: true_label_head = true_label[loss_head] # During eval, even if our model does not predict a NIC candidate, we allow for a NIC gold QID # This qid gets assigned the label of -2 and is always incorrect in eval_wapper # As NLLLoss assumes classes of 0 to #classes-1 except for pad idx, we manually mask # the -2 labels for the loss computation only. As this is just for eval, it won't matter. if not training: label_mask = true_label_head == -2 true_label_head[label_mask] = -1 # batch x M x K -> transpose -> swap K classes with M spans for "k-dimensional" NLLloss log_probs = eval_utils.masked_class_logsoftmax(pred=out, mask=~mask).transpose(1,2) temp = self.crit_pred(log_probs, true_label_head.long().to(device)) loss.add_loss(loss_head, temp) if not training: true_label_head[label_mask] = -2 return loss
def test_masked_class_logsoftmax_grads(self): # check gradients on preds since that will go back into the rest of the network preds = torch.tensor( [[[2.0, 4.0, 1.0], [3.0, 5.0, 4.0], [1.0, 4.0, 6.0]]], requires_grad=True) entity_ids = torch.tensor([[[1, 3, -1], [5, -1, -1], [4, 5, 6]]]) true_entity_class = torch.tensor([[1, 0, 2]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax( pred=preds, mask=mask).transpose(1, 2) pred_loss = torch.nn.NLLLoss(ignore_index=-1)(pred_log_preds, true_entity_class) pred_loss.backward() actual_grad = preds.grad # we want zero grads on masked candidates masked_actual_grad = torch.where(entity_ids > 0, torch.ones_like(preds), actual_grad) # just put 1's where we want non-zeros and use mask above to only compare padded gradients expected_grad = torch.tensor([[[1.0, 1.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 1.0]]]) assert torch.allclose(expected_grad, masked_actual_grad) # we want to match pytorch when NOT using masking # zero out the gradient to call backward again preds.grad.zero_() # no masking now entity_ids = torch.tensor([[[1, 3, 1], [5, 4, 8], [4, 5, 6]]]) true_entity_class = torch.tensor([[1, 0, 2]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax( pred=preds, mask=mask).transpose(1, 2) pred_loss = torch.nn.NLLLoss(ignore_index=-1)(pred_log_preds, true_entity_class) pred_loss.backward() # clone so we can call backward again and zero out the grad actual_grad = preds.grad.clone() preds.grad.zero_() torch_loss_fn = torch.nn.CrossEntropyLoss() torch_loss = torch_loss_fn(preds.transpose(1, 2), true_entity_class) torch_loss.backward() torch_grad = preds.grad assert torch.allclose(torch_grad, actual_grad)
def disambig_output(intermediate_output_dict): """Function to return the probs for a task in Emmental. Args: intermediate_output_dict: output dict from Emmental task flow Returns: NED probabilities for candidates (B x M x K) """ out = intermediate_output_dict[PRED_LAYER]["final_scores"][DISAMBIG][ FINAL_LOSS] mask = intermediate_output_dict["_input_"]["entity_cand_eid_mask"] return eval_utils.masked_class_logsoftmax(pred=out, mask=~mask).exp()
def test_masked_class_logsoftmax_masking(self): preds = torch.tensor([[[2., 4., 1.], [3., 5., 4.]]]) entity_ids = torch.tensor([[[1, 3, -1], [5, -1, -1]]]) first_sample = torch.tensor([[2., 4.]]) denom_0 = torch.log(torch.sum(torch.exp(first_sample))) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) # we only need to match on non-masked values expected_log_probs = torch.tensor([[[ first_sample[0][0]-denom_0, first_sample[0][1]-denom_0, 0], [0, 0, 0]]]) pred_log_preds = eval_utils.masked_class_logsoftmax(pred=preds, mask=mask) * mask assert torch.allclose(expected_log_probs, pred_log_preds)
def test_masked_class_logsoftmax_with_loss(self): # shape batch x M x K # model outputs preds = torch.tensor([[[2., 2., 1.], [3., 5., 4.]]]) # all that matters for this test is that the below is non-negative # since negative indicates masking entity_ids = torch.tensor([[[1, 3, 4], [5, 3, 1]]]) true_entity_class = torch.tensor([[0,1]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax(pred=preds, mask=mask).transpose(1,2) pred_loss = self.trainer.scorer.crit_pred(pred_log_preds, true_entity_class) torch_loss_fn = torch.nn.CrossEntropyLoss() # predictions need to be batch_size x K x M torch_loss = torch_loss_fn(preds.transpose(1,2), true_entity_class) assert torch.allclose(torch_loss, pred_loss)
def test_masked_class_logsoftmax_grads_full_mask(self): preds = torch.tensor([[[2., 4.], [3., 5.], [1., 4.]]], requires_grad=True) # batch x M x K entity_ids = torch.tensor([[[1, -1], [-1, -1], [4, 5]]]) # batch x M true_entity_class = torch.tensor([[0, -1, 1]]) mask = torch.where(entity_ids < 0, torch.zeros_like(preds), torch.ones_like(preds)) pred_log_preds = eval_utils.masked_class_logsoftmax(pred=preds, mask=mask).transpose(1,2) pred_loss = self.trainer.scorer.crit_pred(pred_log_preds, true_entity_class) pred_loss.backward() actual_grad = preds.grad true_entity_class_expanded = true_entity_class.unsqueeze(-1).expand_as(entity_ids) masked_actual_grad = torch.where((entity_ids != -1) & (true_entity_class_expanded != -1), torch.ones_like(preds), actual_grad) # just put 1's where we want non-zeros and use mask above to only compare padded gradients expected_grad = torch.tensor([[[1., 0.], [0., 0.], [1., 1.]]]) assert torch.allclose(expected_grad, masked_actual_grad)
def label_mentions(self, text): sample = self.extract_mentions(text) idxs_arr, aliases_to_predict_per_split, spans_arr, phrase_tokens_arr = sentence_utils.split_sentence( max_aliases=self.args.data_config.max_aliases, phrase=sample['sentence'], spans=sample['spans'], aliases=sample['aliases'], aliases_seen_by_model=[i for i in range(len(sample['aliases']))], seq_len=self.args.data_config.max_word_token_len, word_symbols=self.word_db) aliases_arr = [[sample['aliases'][idx] for idx in idxs] for idxs in idxs_arr] qids_arr = [[sample['qids'][idx] for idx in idxs] for idxs in idxs_arr] word_indices_arr = [ self.word_db.convert_tokens_to_ids(pt) for pt in phrase_tokens_arr ] if len(idxs_arr) > 1: #TODO: support sentences that overflow due to long sequence length or too many mentions raise ValueError( 'Overflowing sentences not currently supported in Annotator') # iterate over each sample in the split for sub_idx in range(len(idxs_arr)): example_aliases = np.ones(self.args.data_config.max_aliases, dtype=np.int) * PAD_ID example_true_entities = np.ones( self.args.data_config.max_aliases) * PAD_ID example_aliases_locs_start = np.ones( self.args.data_config.max_aliases) * PAD_ID example_aliases_locs_end = np.ones( self.args.data_config.max_aliases) * PAD_ID aliases = aliases_arr[sub_idx] for mention_idx, alias in enumerate(aliases): # get aliases alias_trie_idx = self.entity_db.get_alias_idx(alias) alias_qids = np.array(self.entity_db.get_qid_cands(alias)) example_aliases[mention_idx] = alias_trie_idx # alias_idx_pair span_idx = spans_arr[sub_idx][mention_idx] span_start_idx, span_end_idx = span_idx example_aliases_locs_start[mention_idx] = span_start_idx example_aliases_locs_end[mention_idx] = span_end_idx # get word indices word_indices = word_indices_arr[sub_idx] # entity indices from alias table (these are the candidates) entity_indices = self.alias_table(example_aliases) # all CPU embs have to retrieved on the fly batch_on_the_fly_data = {} for emb_name, emb in self.batch_on_the_fly_embs.items(): batch_on_the_fly_data[emb_name] = torch.tensor( emb.batch_prep(example_aliases, entity_indices), device=self.device) outs, entity_pack, _ = self.model( alias_idx_pair_sent=[ torch.tensor(example_aliases_locs_start, device=self.device).unsqueeze(0), torch.tensor(example_aliases_locs_end, device=self.device).unsqueeze(0) ], word_indices=torch.tensor([word_indices], device=self.device), alias_indices=torch.tensor(example_aliases, device=self.device).unsqueeze(0), entity_indices=torch.tensor(entity_indices, device=self.device).unsqueeze(0), batch_prepped_data={}, batch_on_the_fly_data=batch_on_the_fly_data) entity_cands = eval_utils.map_aliases_to_candidates( self.args.data_config.train_in_candidates, self.entity_db, aliases) # recover predictions probs = torch.exp( eval_utils.masked_class_logsoftmax( pred=outs[DISAMBIG][FINAL_LOSS], mask=~entity_pack.mask, dim=2)) max_probs, max_probs_indices = probs.max(2) pred_cands = [] pred_probs = [] titles = [] for alias_idx in range(len(aliases)): pred_idx = max_probs_indices[0][alias_idx] pred_prob = max_probs[0][alias_idx].item() pred_qid = entity_cands[alias_idx][pred_idx] if pred_prob > self.threshold: pred_cands.append(pred_qid) pred_probs.append(pred_prob) titles.append( self.entity_db. get_title(pred_qid) if pred_qid != 'NC' else 'NC') return pred_cands, pred_probs, titles
def forward(self, context_matrix_dict, alias_idx_pair_sent, entity_pack, sent_emb): out = {DISAMBIG: {}, INDICATOR: {}} indicator_outputs = {} expert_slice_repr = {} predictor_outputs = {} if "context_matrix_main" not in context_matrix_dict: context_matrix_dict[ "context_matrix_main"] = model_utils.generate_final_context_matrix( context_matrix_dict) context_matrix = context_matrix_dict["context_matrix_main"] batch_size, M, K, H = context_matrix.shape assert M == self.M assert K == self.K assert H == self.hidden_size for i, slice_head in enumerate(self.train_heads): # Generate slice expert representation per head # context_matrix is batch x M x K x H expert_slice_repr[slice_head] = self.transform_modules[slice_head]( context_matrix) # Pass the expert slice representation through the shared prediction layer # Predictor_outputs is batch x M x K predictor_outputs[slice_head] = self.shared_slice_pred_head( expert_slice_repr[slice_head]).squeeze(-1) # Get an alias_matrix output (batch x M x H) # TODO: remove extra inputs alias_matrix, alias_word_weights = self.ind_alias_mha[slice_head]( sent_embedding=sent_emb, entity_embedding=context_matrix, entity_mask=entity_pack.mask, alias_idx_pair_sent=alias_idx_pair_sent, slice_emb_alias=self.slice_emb_ind_alias( torch.tensor(i, device=context_matrix.device)), slice_emb_ent=self.slice_emb_ind_ent( torch.tensor(i, device=context_matrix.device))) # Get indicator head outputs; size batch x M x 2 per head indicator_outputs[slice_head] = self.indicator_heads[slice_head]( alias_matrix) # Generate predictions via softmax + taking the "positive" class label # Output size is batch x M x num_slices indicator_preds = torch.cat( [ F.log_softmax(indicator_outputs[slice_head], dim=-1)[:, :, 1].unsqueeze(-1) for slice_head in self.train_heads ], dim=-1, ) assert not torch.isnan(indicator_preds).any() assert not torch.isinf(indicator_preds).any() # Compute the "confidence" # Output size should be batch x M x K x num_slices predictor_confidences = torch.cat( [ eval_utils.masked_class_logsoftmax( pred=predictor_outputs[slice_head], mask=~entity_pack.mask).unsqueeze(-1) for slice_head in self.train_heads ], dim=-1, ) assert not torch.isnan(predictor_confidences).any() assert not torch.isinf(predictor_confidences).any() if self.use_ind_attn: prod = indicator_preds # * margin_confidence prod[prod < 0.1] = -100000.0 * self.temperature attention_weights = F.softmax(prod / self.temperature, dim=-1) else: # Take margin confidence over K values to generate confidences of batch x M x num_slices vals, indices = torch.topk(predictor_confidences, k=2, dim=2) margin_confidence = (vals[:, :, 0, :] - vals[:, :, 1, :]) / vals.sum(2) assert list(margin_confidence.shape) == [ batch_size, self.M, len(self.train_heads) ] attention_weights = F.softmax( (indicator_preds + margin_confidence) / self.temperature, dim=-1) assert not torch.isnan(attention_weights).any() assert not torch.isinf(attention_weights).any() # attention_weights is batch x M x num_slices # slice_representations is batch_size x M x K x num_slices x feat_dim slice_representations = torch.stack( [expert_slice_repr[slice_head] for slice_head in self.train_heads], dim=3) # attention_weights becomes batch_size x M x K x num_slices x H of slice_representations attention_weights = attention_weights.unsqueeze(2).unsqueeze( -1).expand_as(slice_representations) # Reweight representations with weighted sum across slices reweighted_rep = torch.sum(attention_weights * slice_representations, dim=3) assert reweighted_rep.shape == context_matrix.shape # Pass through the final prediction layer for slice_head in self.train_heads: out[DISAMBIG][train_utils.get_slice_head_pred_name( slice_head)] = predictor_outputs[slice_head] for slice_head in self.train_heads: out[INDICATOR][train_utils.get_slice_head_ind_name( slice_head)] = indicator_outputs[slice_head] # Used for debugging if self.remove_final_loss: out[DISAMBIG][FINAL_LOSS] = out[DISAMBIG][ train_utils.get_slice_head_pred_name(BASE_SLICE)] else: out[DISAMBIG][FINAL_LOSS] = self.final_pred_head( reweighted_rep).squeeze(-1) return out, reweighted_rep
def forward(self, alias_idx_pair_sent, sent_embedding, entity_embedding, batch_prepped, batch_on_the_fly_data): batch_size = sent_embedding.tensor.shape[0] out = {DISAMBIG: {}} # Prepare inputs for attention modules # Get the KG metadata for the KG module - this is 0/1 if pair is connected if REL_INDICES_KEY in batch_on_the_fly_data: kg_bias = batch_on_the_fly_data[REL_INDICES_KEY].float().to( sent_embedding.tensor.device).reshape(batch_size, self.M * self.K, self.M * self.K) else: kg_bias = torch.zeros(batch_size, self.M * self.K, self.M * self.K).to(sent_embedding.tensor.device) kg_bias_diag = kg_bias + self.kg_weight * torch.eye( self.M * self.K).repeat(batch_size, 1, 1).view( batch_size, self.M * self.K, self.M * self.K).to( kg_bias.device) kg_bias_norm = self.softmax( kg_bias_diag.masked_fill((kg_bias_diag == 0), float(-1e9))) sent_tensor = sent_embedding.tensor.transpose(0, 1) # Resize the alias embeddings and the entity mask from B x M x K x D to B x (M*K) x D entity_mask = entity_embedding.mask entity_embedding = entity_embedding.tensor.contiguous().view( batch_size, self.M * self.K, self.hidden_size) entity_embedding = entity_embedding.transpose( 0, 1) # reshape for attention key_padding_mask_entities = entity_mask.contiguous().view( batch_size, self.M * self.K) # Iterate through stages query_tensor = entity_embedding for stage_index in range(self.num_model_stages): # As we are adding a residual in the attention modules, we can make embs empty embs = [] #============================================================================ # Phrase module: compute attention between entities and words #============================================================================ word_entity_attn_context, word_entity_attn_weights = self.attention_modules[ f"stage_{stage_index}_entity_word"]( q=query_tensor, x=sent_tensor, key_mask=sent_embedding.mask, attn_mask=None) # Add embeddings to be merged in the output embs.append(word_entity_attn_context) # Save the attention weights self.attention_weights[ f"stage_{stage_index}_entity_word"] = word_entity_attn_weights #============================================================================ # Co-occurrence module: compute self attention over entities #============================================================================ # Move entity mask over self.e2e_entity_mask = self.e2e_entity_mask.to( key_padding_mask_entities.device) entity_attn_context, entity_attn_weights = self.attention_modules[ f"stage_{stage_index}_self_entity"]( x=query_tensor, key_mask=key_padding_mask_entities, attn_mask=self.e2e_entity_mask) # Mask out MxK of single aliases, alias_indices is batch x M, mask is true when single alias non_null_aliases = (self.K - key_padding_mask_entities.reshape( batch_size, self.M, self.K).sum(-1)) != 0 entity_attn_post_mask = ( non_null_aliases.sum(1) == 1).unsqueeze(1).expand( batch_size, self.K * self.M).transpose(0, 1) entity_attn_post_mask = entity_attn_post_mask.unsqueeze( -1).expand_as(entity_attn_context) entity_attn_context = torch.where( entity_attn_post_mask, torch.zeros_like(entity_attn_context), entity_attn_context) # Add embeddings to be merged in the output embs.append(entity_attn_context) # Save the attention weights self.attention_weights[ f"stage_{stage_index}_self_entity"] = entity_attn_weights context_matrix = self.combine_modules[ f"stage_{stage_index}_combine"](embs) #============================================================================ # KG module: add in KG connectivity bias #============================================================================ context_matrix_kg = torch.bmm(kg_bias_norm, context_matrix.transpose( 0, 1)).transpose(0, 1) if stage_index < self.num_model_stages - 1: pred = self.predict_layers[DISAMBIG][ train_utils.get_stage_head_name(stage_index)]( context_matrix) pred = pred.transpose(0, 1).squeeze(2).reshape( batch_size, self.M, self.K) pred_kg = self.predict_layers[DISAMBIG][ train_utils.get_stage_head_name(stage_index)]( (context_matrix + context_matrix_kg) / 2) pred_kg = pred_kg.transpose(0, 1).squeeze(2).reshape( batch_size, self.M, self.K) out[DISAMBIG][ f"{train_utils.get_stage_head_name(stage_index)}"] = torch.max( torch.cat([pred.unsqueeze(3), pred_kg.unsqueeze(3)], dim=-1), dim=-1)[0] pred_logit = eval_utils.masked_class_logsoftmax( pred=pred, mask=~entity_mask) context_norm = model_utils.normalize_matrix( (context_matrix + context_matrix_kg) / 2, dim=2) assert not torch.isnan(context_norm).any() assert not torch.isinf(context_norm).any() # Add predictions so model can learn from previous predictions context_matrix = context_norm + pred_logit.contiguous().view( batch_size, self.M * self.K, 1).transpose(0, 1) else: context_matrix_nokg = context_matrix context_matrix = (context_matrix + context_matrix_kg) / 2 context_matrix_main = context_matrix query_tensor = context_matrix context_matrix_nokg = context_matrix_nokg.transpose(0, 1).reshape( batch_size, self.M, self.K, self.hidden_size) context_matrix_main = context_matrix_main.transpose(0, 1).reshape( batch_size, self.M, self.K, self.hidden_size) # context_mat_dict is the contextualized entity embeddings, out is the predictions context_mat_dict = { "context_matrix_nokg": context_matrix_nokg, MAIN_CONTEXT_MATRIX: context_matrix_main } return context_mat_dict, out
def label_mentions(self, text_list): if type(text_list) is str: text_list = [text_list] else: assert type(text_list) is list and len(text_list) > 0 and type( text_list[0]) is str, f"We only accept inputs of strings and lists of strings" ebs = self.args.run_config.eval_batch_size total_start_exs = 0 total_final_exs = 0 dropped_by_thresh = 0 final_char_spans = [] batch_example_aliases = [] batch_example_aliases_locs_start = [] batch_example_aliases_locs_end = [] batch_example_alias_list_pos = [] batch_example_true_entities = [] batch_word_indices = [] batch_spans_arr = [] batch_aliases_arr = [] batch_idx_unq = [] batch_subsplit_idx = [] for idx_unq, text in tqdm(enumerate(text_list), desc="Prepping data", total=len(text_list)): sample = self.extract_mentions(text) total_start_exs += len(sample['aliases']) char_spans = self.get_char_spans(sample['spans'], text) final_char_spans.append(char_spans) idxs_arr, aliases_to_predict_per_split, spans_arr, phrase_tokens_arr, pos_idxs = sentence_utils.split_sentence( max_aliases=self.args.data_config.max_aliases, phrase=sample['sentence'], spans=sample['spans'], aliases=sample['aliases'], aliases_seen_by_model=[i for i in range(len(sample['aliases']))], seq_len=self.args.data_config.max_word_token_len, word_symbols=self.word_db) aliases_arr = [[sample['aliases'][idx] for idx in idxs] for idxs in idxs_arr] old_spans_arr = [[sample['spans'][idx] for idx in idxs] for idxs in idxs_arr] qids_arr = [[sample['qids'][idx] for idx in idxs] for idxs in idxs_arr] word_indices_arr = [self.word_db.convert_tokens_to_ids(pt) for pt in phrase_tokens_arr] # iterate over each sample in the split for sub_idx in range(len(idxs_arr)): # ==================================================== # GENERATE MODEL INPUTS # ==================================================== aliases_to_predict_arr = aliases_to_predict_per_split[sub_idx] assert len(aliases_to_predict_arr) >= 0, f'There are no aliases to predict for an example. This should not happen at this point.' assert len(aliases_arr[ sub_idx]) <= self.args.data_config.max_aliases, f'Each example should have no more that {self.args.data_config.max_aliases} max aliases. {sample} does.' example_aliases = np.ones(self.args.data_config.max_aliases) * PAD_ID example_aliases_locs_start = np.ones(self.args.data_config.max_aliases) * PAD_ID example_aliases_locs_end = np.ones(self.args.data_config.max_aliases) * PAD_ID example_alias_list_pos = np.ones(self.args.data_config.max_aliases) * PAD_ID example_true_entities = np.ones(self.args.data_config.max_aliases) * PAD_ID for mention_idx, alias in enumerate(aliases_arr[sub_idx]): span_start_idx, span_end_idx = spans_arr[sub_idx][mention_idx] # generate indexes into alias table. alias_trie_idx = self.entity_db.get_alias_idx(alias) alias_qids = np.array(self.entity_db.get_qid_cands(alias)) if not qids_arr[sub_idx][mention_idx] in alias_qids: # assert not data_args.train_in_candidates if not self.args.data_config.train_in_candidates: # set class label to be "not in candidate set" true_entity_idx = 0 else: true_entity_idx = -2 else: # Here we are getting the correct class label for training. # Our training is "which of the max_entities entity candidates is the right one (class labels 1 to max_entities) or is it none of these (class label 0)". # + (not discard_noncandidate_entities) is to ensure label 0 is reserved for "not in candidate set" class true_entity_idx = np.nonzero(alias_qids == qids_arr[sub_idx][mention_idx])[0][0] + ( not self.args.data_config.train_in_candidates) example_aliases[mention_idx] = alias_trie_idx example_aliases_locs_start[mention_idx] = span_start_idx # The span_idxs are [start, end). We want [start, end]. So subtract 1 from end idx. example_aliases_locs_end[mention_idx] = span_end_idx - 1 example_alias_list_pos[mention_idx] = idxs_arr[sub_idx][mention_idx] # leave as -1 if it's not an alias we want to predict; we get these if we split a sentence and need to only predict subsets if mention_idx in aliases_to_predict_arr: example_true_entities[mention_idx] = true_entity_idx # get word indices word_indices = word_indices_arr[sub_idx] batch_example_aliases.append(example_aliases) batch_example_aliases_locs_start.append(example_aliases_locs_start) batch_example_aliases_locs_end.append(example_aliases_locs_end) batch_example_alias_list_pos.append(example_alias_list_pos) batch_example_true_entities.append(example_true_entities) batch_word_indices.append(word_indices) batch_aliases_arr.append(aliases_arr[sub_idx]) # Add the orginal sample spans because spans_arr is w.r.t BERT subword token batch_spans_arr.append(old_spans_arr[sub_idx]) batch_idx_unq.append(idx_unq) batch_subsplit_idx.append(sub_idx) batch_example_aliases = torch.tensor(batch_example_aliases).long() batch_example_aliases_locs_start = torch.tensor(batch_example_aliases_locs_start, device=self.device) batch_example_aliases_locs_end = torch.tensor(batch_example_aliases_locs_end, device=self.device) batch_example_true_entities = torch.tensor(batch_example_true_entities, device=self.device) batch_word_indices = torch.tensor(batch_word_indices, device=self.device) final_pred_cands = [[] for _ in range(len(text_list))] final_all_cands = [[] for _ in range(len(text_list))] final_cand_probs = [[] for _ in range(len(text_list))] final_pred_probs = [[] for _ in range(len(text_list))] final_titles = [[] for _ in range(len(text_list))] final_spans = [[] for _ in range(len(text_list))] final_aliases = [[] for _ in range(len(text_list))] for b_i in tqdm(range(0, batch_example_aliases.shape[0], ebs), desc="Evaluating model"): # entity indices from alias table (these are the candidates) batch_entity_indices = self.alias_table(batch_example_aliases[b_i:b_i + ebs]) # all CPU embs have to retrieved on the fly batch_on_the_fly_data = {} for emb_name, emb in self.batch_on_the_fly_embs.items(): batch_prep = [] for j in range(b_i, min(b_i + ebs, batch_example_aliases.shape[0])): batch_prep.append(emb.batch_prep(batch_example_aliases[j], batch_entity_indices[j - b_i])) batch_on_the_fly_data[emb_name] = torch.tensor(batch_prep, device=self.device) alias_idx_pair_sent = [batch_example_aliases_locs_start[b_i:b_i + ebs], batch_example_aliases_locs_end[b_i:b_i + ebs]] word_indices = batch_word_indices[b_i:b_i + ebs] alias_indices = batch_example_aliases[b_i:b_i + ebs] entity_indices = torch.tensor(batch_entity_indices, device=self.device) outs, entity_pack, _ = self.model( alias_idx_pair_sent=alias_idx_pair_sent, word_indices=word_indices, alias_indices=alias_indices, entity_indices=entity_indices, batch_prepped_data={}, batch_on_the_fly_data=batch_on_the_fly_data) # ==================================================== # EVALUATE MODEL OUTPUTS # ==================================================== final_loss_vals = outs[DISAMBIG][FINAL_LOSS] # recover predictions probs = torch.exp(eval_utils.masked_class_logsoftmax(pred=final_loss_vals, mask=~entity_pack.mask, dim=2)) max_probs, max_probs_indices = probs.max(2) for ex_i in range(final_loss_vals.shape[0]): idx_unq = batch_idx_unq[b_i + ex_i] subsplit_idx = batch_subsplit_idx[b_i + ex_i] entity_cands = eval_utils.map_aliases_to_candidates(self.args.data_config.train_in_candidates, self.entity_db, batch_aliases_arr[b_i + ex_i]) # batch size is 1 so we can reshape probs_ex = probs[ex_i].detach().cpu().numpy().reshape(self.args.data_config.max_aliases, probs.shape[2]) for alias_idx, true_entity_pos_idx in enumerate(batch_example_true_entities[b_i + ex_i]): if true_entity_pos_idx != PAD_ID: pred_idx = max_probs_indices[ex_i][alias_idx] pred_prob = max_probs[ex_i][alias_idx].item() all_cands = entity_cands[alias_idx] pred_qid = all_cands[pred_idx] if pred_prob > self.threshold: final_all_cands[idx_unq].append(all_cands) final_cand_probs[idx_unq].append(probs_ex[alias_idx]) final_pred_cands[idx_unq].append(pred_qid) final_pred_probs[idx_unq].append(pred_prob) final_aliases[idx_unq].append(batch_aliases_arr[b_i + ex_i][alias_idx]) final_spans[idx_unq].append(batch_spans_arr[b_i + ex_i][alias_idx]) final_titles[idx_unq].append(self.entity_db.get_title(pred_qid) if pred_qid != 'NC' else 'NC') total_final_exs += 1 else: dropped_by_thresh += 1 assert total_final_exs + dropped_by_thresh == total_start_exs, f"Something went wrong and we have predicted fewer mentions than extracted. Start {total_start_exs}, Out {total_final_exs}, No cand {dropped_by_thresh}" return final_pred_cands, final_pred_probs, final_titles, final_all_cands, final_cand_probs, final_spans, final_aliases