def propagate_message0(self, cmd, x_loc, x_ctx, x_ctx_var_drop, x_mask, adj, types): x_ctx = x_ctx * x_ctx_var_drop proj_x_loc = self.project_x_loc(self.read_drop(x_loc)) proj_x_ctx = self.project_x_ctx(self.read_drop(x_ctx)) x_joint = torch.cat([x_loc, x_ctx, proj_x_loc * proj_x_ctx], dim=-1) queries = self.queries(x_joint) #w6 keys = self.keys(x_joint) * self.proj_keys(cmd)[:, None, :] vals = self.vals(x_joint) * self.proj_vals(cmd)[:, None, :] batch_size = keys.size(0) N = keys.size(1) edge_score = (torch.bmm(queries, torch.transpose(keys, 1, 2)) / np.sqrt(self.node_dim)) edge_score = util.replace_masked_values(edge_score, adj > 0, -1e30) edge_prob = F.softmax(edge_score, dim=-1) edge_prob = util.replace_masked_values(edge_prob, adj > 0, 0) message = torch.bmm(edge_prob, vals) x_ctx_new = self.mem_update(torch.cat([x_ctx, message], dim=-1)) return x_ctx_new
def _get_combined_likelihood(self, answer_as_list_of_bios, log_probs): # answer_as_list_of_bios - Shape: (batch_size, # of correct sequences, seq_length) # log_probs - Shape: (batch_size, seq_length, 3) # Shape: (batch_size, # of correct sequences, seq_length, 3) # duplicate log_probs for each gold bios sequence expanded_log_probs = log_probs.unsqueeze(1).expand( -1, answer_as_list_of_bios.size()[1], -1, -1) # get the log-likelihood per each sequence index # Shape: (batch_size, # of correct sequences, seq_length) log_likelihoods = \ torch.gather(expanded_log_probs, dim=-1, index=answer_as_list_of_bios.unsqueeze(-1)).squeeze(-1) # Shape: (batch_size, # of correct sequences) correct_sequences_pad_mask = (answer_as_list_of_bios.sum(-1) > 0).long() # Sum the log-likelihoods for each index to get the log-likelihood of the sequence # Shape: (batch_size, # of correct sequences) sequences_log_likelihoods = log_likelihoods.sum(dim=-1) sequences_log_likelihoods = replace_masked_values( sequences_log_likelihoods, correct_sequences_pad_mask, -1e7) # Sum the log-likelihoods for each sequence to get the marginalized log-likelihood over the correct answers log_marginal_likelihood = logsumexp(sequences_log_likelihoods, dim=-1) return log_marginal_likelihood
def take_step( self, last_predictions: torch.Tensor, state: Dict[str, torch.Tensor] ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: # get the relevant scores for the time step class_log_probabilities = state['log_probs'][:, state['step_num'][0], :] is_wordpiece = ( 1 - state['wordpiece_mask'][:, state['step_num'][0]]).byte() # mask illegal BIO transitions transitions_mask = torch.cat( (torch.ones_like(class_log_probabilities[:, :3]), torch.zeros_like(class_log_probabilities[:, -2:])), dim=-1).byte() transitions_mask[:, 2] &= ((last_predictions == 1) | (last_predictions == 2)).byte() transitions_mask[:, 1:3] &= ((class_log_probabilities[:, :3] == 0.0).sum(-1) != 3).byte().unsqueeze(-1).repeat(1, 2) # assuming the wordpiece mask doesn't intersect with the other masks (pad, cls/sep) transitions_mask[:, 2] |= is_wordpiece & ((last_predictions == 1) | (last_predictions == 2)).byte() class_log_probabilities = replace_masked_values( class_log_probabilities, transitions_mask, -1e7) state['step_num'] = state['step_num'].clone() + 1 return class_log_probabilities, state
def module(self, bert_out, seq_mask=None): logits = self.predictor(bert_out) if self._use_crf: # The mask should not be applied here when using CRF, but should be passed ot the CRF log_probs = torch.nn.functional.log_softmax(logits, dim=-1) else: if seq_mask is not None: log_probs = replace_masked_values( torch.nn.functional.log_softmax(logits, dim=-1), seq_mask.unsqueeze(-1), 0.0) logits = replace_masked_values(logits, seq_mask.unsqueeze(-1), -1e7) else: log_probs = torch.nn.functional.log_softmax(logits) return log_probs, logits
def log_likelihood(self, gold_labels, log_probs, seq_mask, is_bio_mask, **kwargs): # we only want the log probabilities of the gold labels # what we get is: # log_likelihoods_for_multispan[i,j] == log_probs[i,j, gold_labels[i,j]] log_likelihoods_for_multispan = \ torch.gather(log_probs, dim=-1, index=gold_labels.unsqueeze(-1)).squeeze(-1) # Our marginal likelihood is the sum of all the gold label likelihoods, ignoring the # padding tokens. log_likelihoods_for_multispan = \ replace_masked_values(log_likelihoods_for_multispan, seq_mask, 0.0) log_marginal_likelihood_for_multispan = log_likelihoods_for_multispan.sum( dim=-1) # For questions without spans, we set their log probabilities to be very small negative value log_marginal_likelihood_for_multispan = \ replace_masked_values(log_marginal_likelihood_for_multispan, is_bio_mask, -1e7) return log_marginal_likelihood_for_multispan
def prediction(self, log_probs, logits, qp_tokens, p_text, q_text, seq_mask, wordpiece_mask, use_beam_search): if use_beam_search: top_k_predictions = self._get_top_k_sequences( log_probs.unsqueeze(0), wordpiece_mask.unsqueeze(0), self._prediction_beam_size) predicted_tags = top_k_predictions[0, 0, :] else: predicted_tags = torch.argmax(logits, dim=-1) predicted_tags = replace_masked_values(predicted_tags, seq_mask, 0) return MultiSpanHead.decode_spans_from_tags(predicted_tags, qp_tokens, p_text, q_text)
def log_likelihood(self, gold_labels, log_probs, seq_mask, is_bio_mask, **kwargs): logits = kwargs['logits'] if gold_labels is not None: log_denominator = self.crf._input_likelihood(logits, seq_mask) log_numerator = self.crf._joint_likelihood(logits, gold_labels, seq_mask) log_likelihood = log_numerator - log_denominator log_likelihood = replace_masked_values(log_likelihood, is_bio_mask, -1e7) return log_likelihood
def log_likelihood(self, answer_as_text_to_disjoint_bios, answer_as_list_of_bios, span_bio_labels, log_probs, logits, seq_mask, wordpiece_mask, is_bio_mask, **kwargs): # answer_as_text_to_disjoint_bios - Shape: (batch_size, # of text answers, # of spans a for text answer, seq_length) # answer_as_list_of_bios - Shape: (batch_size, # of correct sequences, seq_length) # log_probs - Shape: (batch_size, seq_length, 3) # seq_mask - Shape: (batch_size, seq_length) # Generate most likely correct predictions if self._use_crf: raise NotImplementedError else: with torch.no_grad(): answer_as_list_of_bios = answer_as_list_of_bios * seq_mask.unsqueeze( 1) if answer_as_text_to_disjoint_bios.sum() > 0: full_bio = span_bio_labels if self._generation_top_k > 0: most_likely_predictions = self._get_top_k_sequences( log_probs, wordpiece_mask, self._generation_top_k) most_likely_predictions = most_likely_predictions * seq_mask.unsqueeze( 1) generated_list_of_bios = self._filter_correct_predictions( most_likely_predictions, answer_as_text_to_disjoint_bios, full_bio) is_pregenerated_answer_format_mask = ( answer_as_list_of_bios.sum((1, 2)) > 0).unsqueeze(-1).unsqueeze(-1).long() list_of_bios = torch.cat( (answer_as_list_of_bios, (generated_list_of_bios * (1 - is_pregenerated_answer_format_mask))), dim=1) list_of_bios = self._add_full_bio( list_of_bios, full_bio) else: is_pregenerated_answer_format_mask = ( answer_as_list_of_bios.sum((1, 2)) > 0).long() list_of_bios = torch.cat( (answer_as_list_of_bios, (full_bio * (1 - is_pregenerated_answer_format_mask ).unsqueeze(-1)).unsqueeze(1)), dim=1) else: list_of_bios = answer_as_list_of_bios ### Calculate log-likelihood from list_of_bios if self._use_crf: raise NotImplementedError else: log_marginal_likelihood_for_multispan = self._get_combined_likelihood( list_of_bios, log_probs) # For questions without spans, we set their log probabilities to be very small negative value log_marginal_likelihood_for_multispan = \ replace_masked_values(log_marginal_likelihood_for_multispan, is_bio_mask, -1e7) return log_marginal_likelihood_for_multispan
def prediction(self, log_probs, logits, qp_tokens, p_text, q_text, mask): predicted_tags = torch.argmax(logits, dim=-1) predicted_tags = replace_masked_values(predicted_tags, mask, 0) return MultiSpanHead.decode_spans_from_tags(predicted_tags, qp_tokens, p_text, q_text)
def forward(self, # type: ignore input_ids: torch.LongTensor, input_mask: torch.LongTensor, input_segments: torch.LongTensor, passage_mask: torch.LongTensor, question_mask: torch.LongTensor, number_indices: torch.LongTensor, passage_number_order: torch.LongTensor, question_number_order: torch.LongTensor, question_number_indices: torch.LongTensor, gnodes: torch.LongTensor, gnodes_len: torch.LongTensor, gnodes_mask: torch.LongTensor, gedges: torch.LongTensor, gedge_types: int, answer_as_passage_spans: torch.LongTensor = None, answer_as_question_spans: torch.LongTensor = None, answer_as_add_sub_expressions: torch.LongTensor = None, answer_as_counts: torch.LongTensor = None, answer_as_text_to_disjoint_bios: torch.LongTensor = None, answer_as_list_of_bios: torch.LongTensor = None, span_bio_labels: torch.LongTensor = None, bio_wordpiece_mask: torch.LongTensor = None, is_bio_mask: torch.LongTensor = None, metadata: List[Dict[str, Any]] = None) -> Dict[str, torch.Tensor]: outputs = self.bert(input_ids, attention_mask=input_mask, token_type_ids=input_segments) sequence_output = outputs[0] sequence_output_list = [ item for item in outputs[2][-4:] ] batch_size = input_ids.size(0) if ("passage_span_extraction" in self.answering_abilities or "question_span" in self.answering_abilities) and self.use_gcn: # M2, M3 sequence_alg = self._gcn_input_proj(torch.cat([sequence_output_list[2], sequence_output_list[3]], dim=2)) encoded_passage_for_numbers = sequence_alg encoded_question_for_numbers = sequence_alg encoded_dates = sequence_alg encoded_passage = sequence_alg #gen embedding real_gnodes_indices = gnodes_mask - 1 real_gnodes_mask = gnodes_mask>0 real_gnodes_indices_col_len = real_gnodes_indices.size(-1) z0 = real_gnodes_indices.view(batch_size, -1) #gnodes will large than 0, since passage_start exist z1=(z0>-9999).nonzero()[:,0].view(batch_size,-1) gnodes_embs = encoded_passage[z1,z0] gnodes_embs = gnodes_embs.view(batch_size,-1,real_gnodes_indices_col_len,encoded_passage.size(-1)) gnodes_token_ids = real_gnodes_indices real_gnodes_mask = real_gnodes_mask.unsqueeze(-1).expand(-1,-1, -1, encoded_passage.size(-1)) gnodes_embs = util.replace_masked_values(gnodes_embs, real_gnodes_mask, 0) sum_gnodes_embs = gnodes_embs.sum(-2) mean_gnodes_embs = sum_gnodes_embs/gnodes_len.unsqueeze(-1).float() gnodes_mask_bin = (gnodes_len > 0).long() cls_emb = sequence_output[:, 0] new_gnodes_embs = self.qdgat(mean_gnodes_embs, gnodes_mask_bin, cls_emb, gedges, gedge_types, is_print=False) gnodes_embs_updated = util.replace_masked_values(new_gnodes_embs.unsqueeze(-2).expand(-1,-1,real_gnodes_indices_col_len,-1), real_gnodes_mask, 0) gnodes_embs_updated = gnodes_embs_updated.view(batch_size, -1, sequence_alg.size(-1)) gcn_info_vec = torch.zeros(encoded_passage.shape, dtype=torch.float, device=gnodes.device) gnodes_updated_index = util.replace_masked_values(z0, z0>0, 0) gcn_info_vec.scatter_(1, gnodes_updated_index.unsqueeze(-1).expand(-1, -1, gnodes_embs_updated.size(-1)), gnodes_embs_updated) sequence_output_list[2] = self._gcn_enc(self._proj_ln(sequence_output_list[2] + gcn_info_vec)) sequence_output_list[0] = self._gcn_enc(self._proj_ln0(sequence_output_list[0] + gcn_info_vec)) sequence_output_list[1] = self._gcn_enc(self._proj_ln1(sequence_output_list[1] + gcn_info_vec)) sequence_output_list[3] = self._gcn_enc(self._proj_ln3(sequence_output_list[3] + gcn_info_vec)) # passage hidden and question hidden sequence_h2_weight = self._proj_sequence_h(sequence_output_list[2]).squeeze(-1) passage_h2_weight = util.masked_softmax(sequence_h2_weight, passage_mask) passage_h2 = util.weighted_sum(sequence_output_list[2], passage_h2_weight) question_h2_weight = util.masked_softmax(sequence_h2_weight, question_mask) question_h2 = util.weighted_sum(sequence_output_list[2], question_h2_weight) # passage g0, g1, g2 question_g0_weight = self._proj_sequence_g0(sequence_output_list[0]).squeeze(-1) question_g0_weight = util.masked_softmax(question_g0_weight, question_mask) question_g0 = util.weighted_sum(sequence_output_list[0], question_g0_weight) question_g1_weight = self._proj_sequence_g1(sequence_output_list[1]).squeeze(-1) question_g1_weight = util.masked_softmax(question_g1_weight, question_mask) question_g1 = util.weighted_sum(sequence_output_list[1], question_g1_weight) question_g2_weight = self._proj_sequence_g2(sequence_output_list[2]).squeeze(-1) question_g2_weight = util.masked_softmax(question_g2_weight, question_mask) question_g2 = util.weighted_sum(sequence_output_list[2], question_g2_weight) if len(self.answering_abilities) > 1: # Shape: (batch_size, number_of_abilities) answer_ability_logits = self._answer_ability_predictor(torch.cat([passage_h2, question_h2, sequence_output[:, 0]], 1)) answer_ability_log_probs = F.log_softmax(answer_ability_logits, -1) best_answer_ability = torch.argmax(answer_ability_log_probs, 1) top_two_answer_abilities = torch.topk(answer_ability_log_probs, k=2, dim=1) real_number_indices = number_indices.squeeze(-1) - 1 number_mask = (real_number_indices > -1).long() clamped_number_indices = util.replace_masked_values(real_number_indices, number_mask, 0) encoded_passage_for_numbers = torch.cat([sequence_output_list[2], sequence_output_list[3]], dim=-1) encoded_numbers = torch.gather(encoded_passage_for_numbers, 1, clamped_number_indices.unsqueeze(-1).expand(-1, -1, encoded_passage_for_numbers.size(-1))) number_weight = self._proj_number(encoded_numbers).squeeze(-1) number_mask = (number_indices > -1).long() number_weight = util.masked_softmax(number_weight, number_mask) number_vector = util.weighted_sum(encoded_numbers, number_weight) if "counting" in self.answering_abilities: # Shape: (batch_size, 10) count_number_logits = self._count_number_predictor(torch.cat([number_vector, passage_h2, question_h2, sequence_output[:, 0]], dim=1)) count_number_log_probs = torch.nn.functional.log_softmax(count_number_logits, -1) # Info about the best count number prediction # Shape: (batch_size,) best_count_number = torch.argmax(count_number_log_probs, -1) best_count_log_prob = torch.gather(count_number_log_probs, 1, best_count_number.unsqueeze(-1)).squeeze(-1) if len(self.answering_abilities) > 1: best_count_log_prob += answer_ability_log_probs[:, self._counting_index] if "passage_span_extraction" in self.answering_abilities or "question_span_extraction" in self.answering_abilities: # start 0, 2 sequence_for_span_start = torch.cat([sequence_output_list[2], sequence_output_list[0], sequence_output_list[2]*question_g2.unsqueeze(1), sequence_output_list[0]*question_g0.unsqueeze(1)], dim=2) sequence_span_start_logits = self._span_start_predictor(sequence_for_span_start).squeeze(-1) # Shape: (batch_size, passage_length, modeling_dim * 2) sequence_for_span_end = torch.cat([sequence_output_list[2], sequence_output_list[1], sequence_output_list[2]*question_g2.unsqueeze(1), sequence_output_list[1]*question_g1.unsqueeze(1)], dim=2) # Shape: (batch_size, passage_length) sequence_span_end_logits = self._span_end_predictor(sequence_for_span_end).squeeze(-1) # Shape: (batch_size, passage_length) if "passage_span_extraction" in self.answering_abilities: passage_span_start_log_probs = util.masked_log_softmax(sequence_span_start_logits, passage_mask) passage_span_end_log_probs = util.masked_log_softmax(sequence_span_end_logits, passage_mask) # Info about the best passage span prediction passage_span_start_logits = util.replace_masked_values(sequence_span_start_logits, passage_mask, -1e7) passage_span_end_logits = util.replace_masked_values(sequence_span_end_logits, passage_mask, -1e7) # Shage: (batch_size, topk, 2) best_passage_span = util.get_best_span(passage_span_start_logits, passage_span_end_logits) if "question_span_extraction" in self.answering_abilities: question_span_start_log_probs = util.masked_log_softmax(sequence_span_start_logits, question_mask) question_span_end_log_probs = util.masked_log_softmax(sequence_span_end_logits, question_mask) # Info about the best question span prediction question_span_start_logits = util.replace_masked_values(sequence_span_start_logits, question_mask, -1e7) question_span_end_logits = util.replace_masked_values(sequence_span_end_logits, question_mask, -1e7) # Shape: (batch_size, topk, 2) best_question_span = util.get_best_span(question_span_start_logits, question_span_end_logits) if "addition_subtraction" in self.answering_abilities: alg_encoded_numbers = torch.cat( [encoded_numbers, question_h2.unsqueeze(1).repeat(1, encoded_numbers.size(1), 1), passage_h2.unsqueeze(1).repeat(1, encoded_numbers.size(1), 1), sequence_output[:, 0].unsqueeze(1).repeat(1, encoded_numbers.size(1), 1) ], 2) # Shape: (batch_size, # of numbers in the passage, 3) number_sign_logits = self._number_sign_predictor(alg_encoded_numbers) number_sign_log_probs = torch.nn.functional.log_softmax(number_sign_logits, -1) # Shape: (batch_size, # of numbers in passage). best_signs_for_numbers = torch.argmax(number_sign_log_probs, -1) # For padding numbers, the best sign masked as 0 (not included). best_signs_for_numbers = util.replace_masked_values(best_signs_for_numbers, number_mask, 0) # Shape: (batch_size, # of numbers in passage) best_signs_log_probs = torch.gather(number_sign_log_probs, 2, best_signs_for_numbers.unsqueeze(-1)).squeeze( -1) # the probs of the masked positions should be 1 so that it will not affect the joint probability # TODO: this is not quite right, since if there are many numbers in the passage, # TODO: the joint probability would be very small. best_signs_log_probs = util.replace_masked_values(best_signs_log_probs, number_mask, 0) # Shape: (batch_size,) best_combination_log_prob = best_signs_log_probs.sum(-1) if len(self.answering_abilities) > 1: best_combination_log_prob += answer_ability_log_probs[:, self._addition_subtraction_index] # add multiple span prediction if bio_wordpiece_mask is None or not self.multispan_use_bio_wordpiece_mask: multispan_mask = input_mask else: multispan_mask = input_mask * bio_wordpiece_mask if "multiple_spans" in self.answering_abilities: if self.multispan_head_name == "flexible_loss": multispan_log_probs, multispan_logits = self._multispan_module(sequence_output, seq_mask=multispan_mask) else: multispan_log_probs, multispan_logits = self._multispan_module(sequence_output) output_dict = {} # If answer is given, compute the loss. if answer_as_passage_spans is not None or answer_as_question_spans is not None or answer_as_add_sub_expressions is not None or answer_as_counts is not None: log_marginal_likelihood_list = [] for answering_ability in self.answering_abilities: if answering_ability == "passage_span_extraction": # Shape: (batch_size, # of answer spans) gold_passage_span_starts = answer_as_passage_spans[:, :, 0] gold_passage_span_ends = answer_as_passage_spans[:, :, 1] # Some spans are padded with index -1, # so we clamp those paddings to 0 and then mask after `torch.gather()`. gold_passage_span_mask = (gold_passage_span_starts != -1).long() clamped_gold_passage_span_starts = util.replace_masked_values(gold_passage_span_starts, gold_passage_span_mask, 0) clamped_gold_passage_span_ends = util.replace_masked_values(gold_passage_span_ends, gold_passage_span_mask, 0) # Shape: (batch_size, # of answer spans) log_likelihood_for_passage_span_starts = torch.gather(passage_span_start_log_probs, 1, clamped_gold_passage_span_starts) log_likelihood_for_passage_span_ends = torch.gather(passage_span_end_log_probs, 1, clamped_gold_passage_span_ends) # Shape: (batch_size, # of answer spans) log_likelihood_for_passage_spans = log_likelihood_for_passage_span_starts + log_likelihood_for_passage_span_ends # For those padded spans, we set their log probabilities to be very small negative value log_likelihood_for_passage_spans = util.replace_masked_values(log_likelihood_for_passage_spans, gold_passage_span_mask, -1e7) # Shape: (batch_size, ) log_marginal_likelihood_for_passage_span = util.logsumexp(log_likelihood_for_passage_spans) # print('passage_span_extraction: ', log_marginal_likelihood_for_passage_span) log_marginal_likelihood_list.append(log_marginal_likelihood_for_passage_span) elif answering_ability == "question_span_extraction": # Shape: (batch_size, # of answer spans) gold_question_span_starts = answer_as_question_spans[:, :, 0] gold_question_span_ends = answer_as_question_spans[:, :, 1] # Some spans are padded with index -1, # so we clamp those paddings to 0 and then mask after `torch.gather()`. gold_question_span_mask = (gold_question_span_starts != -1).long() clamped_gold_question_span_starts = util.replace_masked_values(gold_question_span_starts, gold_question_span_mask, 0) clamped_gold_question_span_ends = util.replace_masked_values(gold_question_span_ends, gold_question_span_mask, 0) # Shape: (batch_size, # of answer spans) log_likelihood_for_question_span_starts = torch.gather(question_span_start_log_probs, 1, clamped_gold_question_span_starts) log_likelihood_for_question_span_ends = torch.gather(question_span_end_log_probs, 1, clamped_gold_question_span_ends) # Shape: (batch_size, # of answer spans) log_likelihood_for_question_spans = log_likelihood_for_question_span_starts + log_likelihood_for_question_span_ends # For those padded spans, we set their log probabilities to be very small negative value log_likelihood_for_question_spans = util.replace_masked_values(log_likelihood_for_question_spans, gold_question_span_mask, -1e7) # Shape: (batch_size, ) # pylint: disable=invalid-name log_marginal_likelihood_for_question_span = util.logsumexp(log_likelihood_for_question_spans) # question multi span prediction log_marginal_likelihood_list.append(log_marginal_likelihood_for_question_span) # print('log_marginal_likelihood_for_question_span: ', log_marginal_likelihood_for_question_span) elif answering_ability == "addition_subtraction": # The padded add-sub combinations use 0 as the signs for all numbers, and we mask them here. # Shape: (batch_size, # of combinations) gold_add_sub_mask = (answer_as_add_sub_expressions.sum(-1) > 0).float() # Shape: (batch_size, # of numbers in the passage, # of combinations) gold_add_sub_signs = answer_as_add_sub_expressions.transpose(1, 2) # Shape: (batch_size, # of numbers in the passage, # of combinations) log_likelihood_for_number_signs = torch.gather(number_sign_log_probs, 2, gold_add_sub_signs) # the log likelihood of the masked positions should be 0 # so that it will not affect the joint probability log_likelihood_for_number_signs = util.replace_masked_values(log_likelihood_for_number_signs, number_mask.unsqueeze(-1), 0) # Shape: (batch_size, # of combinations) log_likelihood_for_add_subs = log_likelihood_for_number_signs.sum(1) # For those padded combinations, we set their log probabilities to be very small negative value log_likelihood_for_add_subs = util.replace_masked_values(log_likelihood_for_add_subs, gold_add_sub_mask, -1e7) # Shape: (batch_size, ) log_marginal_likelihood_for_add_sub = util.logsumexp(log_likelihood_for_add_subs) log_marginal_likelihood_list.append(log_marginal_likelihood_for_add_sub) # print('log_marginal_likelihood_for_add_sub: ', log_marginal_likelihood_for_add_sub) elif answering_ability == "counting": # Count answers are padded with label -1, # so we clamp those paddings to 0 and then mask after `torch.gather()`. # Shape: (batch_size, # of count answers) gold_count_mask = (answer_as_counts != -1).long() # Shape: (batch_size, # of count answers) clamped_gold_counts = util.replace_masked_values(answer_as_counts, gold_count_mask, 0) log_likelihood_for_counts = torch.gather(count_number_log_probs, 1, clamped_gold_counts) # For those padded spans, we set their log probabilities to be very small negative value log_likelihood_for_counts = util.replace_masked_values(log_likelihood_for_counts, gold_count_mask, -1e7) # Shape: (batch_size, ) log_marginal_likelihood_for_count = util.logsumexp(log_likelihood_for_counts) log_marginal_likelihood_list.append(log_marginal_likelihood_for_count) # print('log_marginal_likelihood_for_count: ', log_marginal_likelihood_for_count) elif answering_ability == "multiple_spans": if self.multispan_head_name == "flexible_loss": log_marginal_likelihood_for_multispan = \ self._multispan_log_likelihood(answer_as_text_to_disjoint_bios, answer_as_list_of_bios, span_bio_labels, multispan_log_probs, multispan_logits, multispan_mask, bio_wordpiece_mask, is_bio_mask) else: log_marginal_likelihood_for_multispan = \ self._multispan_log_likelihood(span_bio_labels, multispan_log_probs, multispan_mask, is_bio_mask, logits=multispan_logits) log_marginal_likelihood_list.append(log_marginal_likelihood_for_multispan) else: raise ValueError(f"Unsupported answering ability: {answering_ability}") if len(self.answering_abilities) > 1: # Add the ability probabilities if there are more than one abilities all_log_marginal_likelihoods = torch.stack(log_marginal_likelihood_list, dim=-1) all_log_marginal_likelihoods = all_log_marginal_likelihoods + answer_ability_log_probs marginal_log_likelihood = util.logsumexp(all_log_marginal_likelihoods) else: marginal_log_likelihood = log_marginal_likelihood_list[0] output_dict["loss"] = - marginal_log_likelihood.mean() # print('batch_loss: ', output_dict["loss"]) with torch.no_grad(): best_answer_ability = best_answer_ability.detach().cpu().numpy() if metadata is not None: output_dict["question_id"] = [] output_dict["answer"] = [] i = 0 while i < batch_size: if len(self.answering_abilities) > 1: answer_index = best_answer_ability[i] predicted_ability_str = self.answering_abilities[answer_index] else: predicted_ability_str = self.answering_abilities[0] answer_json: Dict[str, Any] = {} question_start = 1 passage_start = len(metadata[i]["question_tokens"]) + 2 # We did not consider multi-mention answers here if predicted_ability_str == "passage_span_extraction": answer_json["answer_type"] = "passage_span" passage_str = metadata[i]['original_passage'] offsets = metadata[i]['passage_token_offsets'] predicted_span = tuple(best_passage_span[i].detach().cpu().numpy()) start_offset = offsets[predicted_span[0] - passage_start][0] end_offset = offsets[predicted_span[1] - passage_start][1] end_offset = end_offset predicted_answer = passage_str[start_offset:end_offset] answer_json["value"] = predicted_answer answer_json["spans"] = [(start_offset, end_offset)] elif predicted_ability_str == "question_span_extraction": answer_json["answer_type"] = "question_span" question_str = metadata[i]['original_question'] offsets = metadata[i]['question_token_offsets'] predicted_span = tuple(best_question_span[i].detach().cpu().numpy()) start_offset = offsets[predicted_span[0] - question_start][0] end_offset = offsets[predicted_span[1] - question_start][1] end_offset = end_offset predicted_answer = question_str[start_offset:end_offset] answer_json["value"] = predicted_answer answer_json["spans"] = [(start_offset, end_offset)] elif predicted_ability_str == "addition_subtraction": # plus_minus combination answer answer_json["answer_type"] = "arithmetic" original_numbers = metadata[i]['original_numbers'] sign_remap = {0: 0, 1: 1, 2: -1} predicted_signs = [sign_remap[it] for it in best_signs_for_numbers[i].detach().cpu().numpy()] result = sum([sign * number for sign, number in zip(predicted_signs, original_numbers)]) predicted_answer = format_number(result) offsets = metadata[i]['passage_token_offsets'] number_indices = metadata[i]['number_indices'] number_positions = [offsets[index - 1] for index in number_indices] answer_json['numbers'] = [] for offset, value, sign in zip(number_positions, original_numbers, predicted_signs): answer_json['numbers'].append({'span': offset, 'value': value, 'sign': sign}) if number_indices[-1] == -1: # There is a dummy 0 number at position -1 added in some cases; we are # removing that here. answer_json["numbers"].pop() answer_json["value"] = result answer_json['number_sign_log_probs'] = number_sign_log_probs[i, :, :].detach().cpu().numpy() elif predicted_ability_str == "counting": answer_json["answer_type"] = "count" predicted_count = best_count_number[i].detach().cpu().numpy() predicted_answer = str(predicted_count) answer_json["count"] = predicted_count elif predicted_ability_str == "multiple_spans": passage_str = metadata[i]["original_passage"] question_str = metadata[i]['original_question'] qp_tokens = metadata[i]["question_passage_tokens"] answer_json["answer_type"] = "multiple_spans" if self.multispan_head_name == "flexible_loss": answer_json["value"], answer_json["spans"], invalid_spans = \ self._multispan_prediction(multispan_log_probs[i], multispan_logits[i], qp_tokens, passage_str, question_str, multispan_mask[i], bio_wordpiece_mask[i], self.multispan_use_prediction_beam_search and not self.training) else: answer_json["value"], answer_json["spans"], invalid_spans = \ self._multispan_prediction(multispan_log_probs[i], multispan_logits[i], qp_tokens, passage_str, question_str, multispan_mask[i]) if self._unique_on_multispan: answer_json["value"] = list(OrderedDict.fromkeys(answer_json["value"])) if self._dont_add_substrings_to_ms: answer_json["value"] = remove_substring_from_prediction(answer_json["value"]) if len(answer_json["value"]) == 0: best_answer_ability[i] = top_two_answer_abilities[1][i][1] continue predicted_answer = answer_json["value"] else: raise ValueError(f"Unsupported answer ability: {predicted_ability_str}") answer_json["predicted_answer"] = predicted_answer output_dict["question_id"].append(metadata[i]["question_id"]) output_dict["answer"].append(answer_json) answer_annotations = metadata[i].get('answer_annotations', []) if answer_annotations: self._drop_metrics(predicted_answer, answer_annotations) ################################################## real_answer_types=[] answer_json['real']={} if answer_as_passage_spans[i][0][0].detach().cpu().numpy() >= 0:#found passage span real_answer_types.append('passage span') if answer_as_question_spans[i][0][0].detach().cpu().numpy() >= 0:#found question span real_answer_types.append('question span') if answer_as_add_sub_expressions[i].sum().detach().cpu().numpy() > 0:#found expr real_answer_types.append('expr') expr_arr = answer_as_add_sub_expressions[i].detach().cpu().numpy() sign_remap = {0: 0, 1: '', 2: '-'} exprs = [] for j in range(min(5,len(expr_arr))): parts=[] for k in range(len(metadata[i]['original_numbers'])): if expr_arr[j][k] > 0: parts.append(str(sign_remap[expr_arr[j][k]])+str(metadata[i]['original_numbers'][k])) exprs.append('+'.join(parts).replace('+-','-',1101)) answer_json['real']['exprs']=';'.join(exprs) if answer_as_counts[i].detach().cpu().numpy()>0: real_answer_types.append('count') answer_json['real']['answer_types']=','.join(real_answer_types) ################################################## i += 1 return output_dict