import torch import torch.nn as nn from model.mention_proposal import MentionProposalModule from model.utils import get_gt_actions from model.memory.entity_memory import EntityMemory from torch.profiler import profile, record_function, ProfilerActivity from typing import Dict, List, Tuple from omegaconf import DictConfig from torch import Tensor from transformers import PreTrainedTokenizerFast import logging import random from collections import defaultdict import copy import time logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO) logger = logging.getLogger() class EntityRankingModel(nn.Module): """ Coreference model based on Entity-Ranking paradigm. In the entity-ranking paradigm, given a new mention we rank the different entity clusters to determine the clustering updates. Entity-Ranking paradigm allows for a naturally scalable solution to coreference resolution. Reference: Rahman and Ng [https://arxiv.org/pdf/1405.5202.pdf] This particular implementation represents the entities/clusters via fixed-dimensional dense representations, typically a simple avereage of mention representations. Clustering is performed in an online, autoregressive manner where mentions are processed in a left-to-right manner. References: Toshniwal et al [https://arxiv.org/pdf/2010.02807.pdf] Toshniwal et al [https://arxiv.org/pdf/2109.09667.pdf] """ def __init__(self, model_config: DictConfig, train_config: DictConfig): super(EntityRankingModel, self).__init__() self.config = model_config self.train_config = train_config # Dropout module - Used during training self.drop_module = nn.Dropout(p=train_config.dropout_rate) self.loss_template_dict = { "total": torch.tensor(0.0, requires_grad=True), "ment_loss": torch.tensor(0.0), "coref": torch.tensor(0.0), "mention_count": torch.tensor(0.0), "ment_correct": torch.tensor(0.001), "ment_total": torch.tensor(0.001), "ment_tp": torch.tensor(0.001), "ment_pp": torch.tensor(0.001), "ment_ap": torch.tensor(0.001), } # Document encoder + Mention proposer self.mention_proposer = MentionProposalModule( self.config, train_config, drop_module=self.drop_module ) # Clustering module span_emb_size: int = self.mention_proposer.span_emb_size # Use of genre feature in clustering or not if self.config.metadata_params.use_genre_feature: self.config.memory.num_feats = 3 self.mem_type = self.config.memory.mem_type.name self.memory_net = EntityMemory( config=self.config.memory, span_emb_size=span_emb_size, drop_module=self.drop_module, ) self.loss_fn = nn.CrossEntropyLoss( label_smoothing=self.train_config.label_smoothing_wt ) if self.config.metadata_params.use_genre_feature: self.genre_embeddings = nn.Embedding( num_embeddings=len(self.config.metadata_params.genres), embedding_dim=self.config.mention_params.emb_size, ) @property def device(self) -> torch.device: return self.mention_proposer.device def get_params(self, named=False) -> Tuple[List, List]: """Returns a tuple of document encoder parameters and rest of the model params.""" encoder_params, mem_params = [], [] for name, param in self.named_parameters(): elem = (name, param) if named else param if "doc_encoder" in name: encoder_params.append(elem) else: mem_params.append(elem) return encoder_params, mem_params def get_tokenizer(self) -> PreTrainedTokenizerFast: """Returns tokenizer used by the document encoder.""" return self.mention_proposer.doc_encoder.get_tokenizer() def get_metadata(self, document: Dict) -> Dict: """Extract metadata such as document genre from document.""" meta_params = self.config.metadata_params if meta_params.use_genre_feature: doc_class = document["doc_key"][:2] if doc_class in meta_params.genres: doc_class_idx = meta_params.genres.index(doc_class) else: doc_class_idx = meta_params.genres.index( meta_params.default_genre ) # Default genre return { "genre": self.genre_embeddings( torch.tensor(doc_class_idx, device=self.device) ) } else: return {} def calculate_coref_loss( self, action_prob_list: List, action_tuple_list: List[Tuple[int, str]] ) -> Tensor: """Calculates the coreference loss for the autoregressive online clustering module. Args: action_prob_list (List): Probability of each clustering action i.e. mention is merged with existing clusters or a new cluster is created. action_tuple_list (List[Tuple[int, str]]): Ground truth actions represented as a tuple of cluster index and action string. 'c' represents that the mention is coreferent with existing clusters while 'o' represents that the mention represents a new cluster. Returns: coref_loss (torch.Tensor): The scalar tensor representing the coreference loss. """ counter = 0 correct = 0 coref_loss = torch.tensor(0.0, device=self.device) num_predictions_clusters = defaultdict(int) for idx, (cell_idx, action_str) in enumerate(action_tuple_list): if action_str == "c": ## Major Entity gt_idx = cell_idx elif action_str == "o": ## Other Entity gt_idx = action_prob_list[counter].shape[0] - 1 else: continue target = torch.tensor([gt_idx], device=self.device) if target[0] == torch.argmax( torch.unsqueeze(action_prob_list[counter], dim=0) ): correct += 1 num_predictions_clusters[ torch.argmax(torch.unsqueeze(action_prob_list[counter], dim=0)).item() ] += 1 coref_loss += self.loss_fn( torch.unsqueeze(action_prob_list[counter], dim=0), target ) counter += 1 return coref_loss @staticmethod def get_filtered_clusters( clusters, init_token_offset, final_token_offset, cluster_mask=None, with_offset=True, ): """Filter clusters from a document given the token offsets.""" """Note that len(cluster_mask) == len(clusters) assured in the previous function.""" filt_clusters = [] no_rep_cluster_mentions = ( [] ) ## Mentions that belonged to a major entity whose representative phrase is not part of the current mentions. for cluster_ind, orig_cluster in enumerate(clusters): cluster = [] for ment_start, ment_end in orig_cluster: if ment_start >= init_token_offset and ment_end < final_token_offset: if with_offset: cluster.append((ment_start, ment_end)) else: cluster.append( ( ment_start - init_token_offset, ment_end - init_token_offset, ) ) if len(cluster) != 0: if ( cluster_mask ): ## During this process if we missed any representative phrases, all clusters that have no representative phrase will be added to the last cluster. if ( cluster_mask[cluster_ind] == True ): ## If representative phrase is in the current segment then, there exists atleast one mention that belongs to the cluster. But anyways filt_clusters.append(cluster) else: no_rep_cluster_mentions.extend(cluster) else: filt_clusters.append(cluster) if cluster_mask: if len(filt_clusters) == 0: filt_clusters.append(no_rep_cluster_mentions) else: filt_clusters[-1].extend(no_rep_cluster_mentions) return filt_clusters @staticmethod def get_filtered_representatives( representatives, init_token_offset, final_token_offset, with_offset=True ): """Filter clusters from a document given the token offsets.""" filt_reps = [] indices = [] for rep_ind, (ment_start, ment_end) in enumerate(representatives): if ment_start >= init_token_offset and ment_end < final_token_offset: if with_offset: filt_reps.append((ment_start, ment_end)) else: filt_reps.append( ( ment_start - init_token_offset, ment_end - init_token_offset, ) ) indices.append(rep_ind) return filt_reps, indices @staticmethod def mask_representative_phrases(rep_emb_list): positive_inds = [] for rep_emb_ind, rep_emb in enumerate(rep_emb_list): if not isinstance(rep_emb, int): positive_inds.append(rep_emb_ind) if len(positive_inds) > 1: num_entitites_preserved = random.randint(1, len(positive_inds)) random.shuffle(positive_inds) for ind in positive_inds[num_entitites_preserved:]: rep_emb_list[ind] = -1 return rep_emb_list def forward_training(self, document: Dict) -> Dict: """Forward pass for training. Args: document: The tensorized document. Returns: loss_dict (Dict): Loss dictionary containing the losses of different stages of the model. """ # print(document["doc_key"]) assert ( len(document["clusters"]) == len(document["representatives"]) + 1 ), "Length of clusters not equal to length of representatives + 1." assert document["representatives"] == sorted( document["representatives"] ), "Representatives are not sorted." loss_dict = copy.deepcopy(self.loss_template_dict) max_training_segments = self.train_config.get("max_training_segments", None) num_segments = len(document["sentences"]) if max_training_segments is None: seg_range = [0, num_segments] else: if num_segments > max_training_segments: start_seg = random.randint(0, num_segments - max_training_segments) seg_range = [start_seg, start_seg + max_training_segments] else: seg_range = [0, num_segments] # Initialize lists to track all the mentions predicted across the chunks pred_mentions_list, mention_emb_list, rep_emb_list = ( [], [], [-1 for _ in range(len(document["representatives"]))], ) init_token_offset = sum( [len(document["sentences"][idx]) for idx in range(0, seg_range[0])] ) token_offset = init_token_offset # Metadata such as document genre can be used by model for clustering metadata = self.get_metadata(document) # Initialize the mention loss ment_loss = None # Step 1: Predict all the mentions for idx in range(seg_range[0], seg_range[1]): num_tokens = len(document["sentences"][idx]) representatives_entities, rep_filtered_inds = ( self.get_filtered_representatives( document["representatives"], token_offset, token_offset + num_tokens, with_offset=False, ) ) cur_doc_slice = { "tensorized_sent": document["tensorized_sent"][idx], "sentence_map": document["sentence_map"][ token_offset : token_offset + num_tokens ], "subtoken_map": document["subtoken_map"][ token_offset : token_offset + num_tokens ], "sent_len_list": [document["sent_len_list"][idx]], "clusters": self.get_filtered_clusters( document["clusters"], token_offset, token_offset + num_tokens, with_offset=False, ), "representatives": representatives_entities, "doc_key": document["doc_key"], } ## No golden mentions in the current segment and mode is golden so basically no job to do. if ( len(cur_doc_slice["clusters"]) == 0 and self.mention_proposer.config.mention_params.use_gold_ments ): token_offset += num_tokens continue proposer_output_dict = self.mention_proposer(cur_doc_slice, eval_loss=True) ### Shifted above because if the model predicts no mentions then earlier it had no mention loss. But now it has. if "ment_loss" in proposer_output_dict: if ment_loss is None: ment_loss = proposer_output_dict["ment_loss"] else: ment_loss += proposer_output_dict["ment_loss"] ## If no mentions are predicted, originally then no coref loss and surprisingly no mention loss as well :) if proposer_output_dict.get("ments", None) is None: token_offset += num_tokens continue ## Mention post-processing and collection happens here: Add the document offset to mentions predicted for the current chunk cur_pred_mentions = proposer_output_dict.get("ments") + token_offset pred_mentions_list.extend(cur_pred_mentions.tolist()) mention_emb_list.extend(proposer_output_dict["ment_emb_list"]) for key in ["ment_correct", "ment_total", "ment_tp", "ment_pp", "ment_ap"]: if key in proposer_output_dict: loss_dict[key] += proposer_output_dict[key] ## Collect representation embeddings: for ind, rep_ind in enumerate(rep_filtered_inds): rep_emb_list[rep_ind] = proposer_output_dict["rep_emb_list"][ind] # Update the document offset for next iteration token_offset += num_tokens ## Collect mention detection loss if ment_loss is not None: ## Tried training the model with only mention loss, but it did not work well. if self.train_config.ment_loss_incl: loss_dict["total"] = ment_loss loss_dict["ment_loss"] = ment_loss # Step 2: Perform clustering # Get clusters part of the truncated document ## select certain entities or representatives if self.train_config.get("generalise", False): rep_emb_list = self.mask_representative_phrases(rep_emb_list) rep_emb_list_filtered = [] entities_mask = [] for rep_emb in rep_emb_list: if not isinstance(rep_emb, int): rep_emb_list_filtered.append(rep_emb) entities_mask.append(True) else: entities_mask.append(False) ## For the other cluster that contains all the mentions that do not belong to any representative phrase. entities_mask.append(True) truncated_document_clusters = { "clusters": self.get_filtered_clusters( document["clusters"], init_token_offset, token_offset, cluster_mask=entities_mask, ) } assert ( len(document["clusters"]) == len(document["representatives"]) + 1 ), "Number of clusters and representatives after segmentation do not match." # Get ground truth clustering mentions gt_actions: List[Tuple[int, str]] = get_gt_actions( pred_mentions_list, truncated_document_clusters, self.config.memory.mem_type ) pred_mentions = torch.tensor(pred_mentions_list, device=self.device) if ( len(rep_emb_list_filtered) == 0 ): ## No representative phrases in the current segments, so no coref loss return loss_dict coref_new_list = self.memory_net.forward_training( pred_mentions, mention_emb_list, rep_emb_list_filtered, gt_actions, metadata ) if len(coref_new_list) > 0: coref_loss = self.calculate_coref_loss(coref_new_list, gt_actions) loss_dict["total"] = loss_dict["total"] + coref_loss loss_dict["coref"] = coref_loss loss_dict["mention_count"] += torch.tensor(len(coref_new_list)) return loss_dict def forward(self, document: Dict, teacher_force=False, gold_mentions=False): """Forward pass of the streaming coreference model. This method performs streaming coreference. The entity clusters from previous documents chunks are represented as vectors and passed along to the processing of subsequent chunks along with the metadata associated with these clusters. Args: document (Dict): Tensorized document Returns: pred_mentions_list (List): Mentions predicted by the mention proposal module mention_scores (List): Scores assigned by the mention proposal module for the predicted mentions gt_actions (List): Ground truth clustering actions; useful for calculating oracle performance action_list (List): Actions predicted by the clustering module for the predicted mentions '""" # Initialize lists to track all the actions taken, mentions predicted across the chunks assert document["representatives"] == sorted( document["representatives"] ), "Representatives are not sorted." print("Device: ", self.device) print("#" * 40) pred_mentions_list, pred_mention_emb_list, mention_scores, pred_actions = ( [], [], [], [], ) # Initialize entity clusters and current document token offset entity_cluster_states, token_offset = None, 0 metadata = self.get_metadata(document) coref_scores_doc = [] link_time = 0.0 for idx in range(0, len(document["sentences"])): num_tokens = len(document["sentences"][idx]) new_representatives_entities, rep_filtered_inds = ( self.get_filtered_representatives( document["representatives"], token_offset, token_offset + num_tokens, with_offset=False, ) ) ext_predicted_mentions_filt, _ = self.get_filtered_representatives( document.get("ext_predicted_mentions", []), token_offset, token_offset + num_tokens, with_offset=False, ) cur_example = { "tensorized_sent": document["tensorized_sent"][idx], "sentence_map": document["sentence_map"][ token_offset : token_offset + num_tokens ], "subtoken_map": document["subtoken_map"][ token_offset : token_offset + num_tokens ], "sent_len_list": [document["sent_len_list"][idx]], "clusters": self.get_filtered_clusters( document["clusters"], token_offset, token_offset + num_tokens, with_offset=False, ), "representatives": new_representatives_entities, "ext_predicted_mentions": ext_predicted_mentions_filt, } # Pass along other metadata for key in document: if key not in cur_example: cur_example[key] = document[key] if len(cur_example["clusters"]) == 0 and ( self.mention_proposer.config.mention_params.use_gold_ments or gold_mentions ): token_offset += num_tokens continue proposer_output_dict = self.mention_proposer( cur_example, gold_mentions=gold_mentions ) if proposer_output_dict.get("ments", None) is None: token_offset += num_tokens continue # Add the document offset to mentions predicted for the current chunk # It's important to add the offset before clustering because features like # number of tokens between the last mention of the cluster and the current mention # will be affected if the current token indices of the mention are not supplied. cur_pred_mentions = proposer_output_dict.get("ments") + token_offset # Update the document offset for next iteration token_offset += num_tokens # Get ground truth clustering mentions pred_mentions_list.extend(cur_pred_mentions.tolist()) gt_actions_full: List[Tuple[int, str]] = get_gt_actions( pred_mentions_list, document, self.config.memory.mem_type ) gt_actions = gt_actions_full[-len(cur_pred_mentions.tolist()) :] pred_mention_emb_list.extend( [emb.tolist() for emb in proposer_output_dict.get("ment_emb_list")] ) mention_scores.extend(proposer_output_dict["ment_scores"].tolist()) start_time = time.time() repr_candidates = list(proposer_output_dict["rep_emb_list"]) # Pass along entity clusters from previous chunks while processing next chunks cur_pred_actions, entity_cluster_states, coref_scores_list = ( self.memory_net( cur_pred_mentions, list(proposer_output_dict["ment_emb_list"]), repr_candidates, gt_actions, metadata, teacher_force=teacher_force, memory_init=entity_cluster_states, ) ) link_time += time.time() - start_time # print( # "Number of representatives available now: ", # entity_cluster_states["mem"].shape[0], # ) pred_actions.extend(cur_pred_actions) coref_scores_doc.extend(coref_scores_list) gt_actions = get_gt_actions( pred_mentions_list, document, self.config.memory.mem_type ) # Useful for oracle calcs for ind in range(len(coref_scores_doc)): coref_scores_doc[ind] = coref_scores_doc[ind].tolist() if entity_cluster_states is not None: for key in entity_cluster_states: entity_cluster_states[key] = entity_cluster_states[key].tolist() return ( pred_mentions_list, pred_mention_emb_list, mention_scores, gt_actions, pred_actions, coref_scores_doc, entity_cluster_states, link_time, )