text: import torch
  import torch.nn as nn
  from transformers import AutoModel
  import util
  import logging
  from collections import Iterable, defaultdict
  import numpy as np
  import torch.nn.init as init
  import higher_order as ho
  from entities import IncrementalEntities, GoldLabelStrategy
  from torch import Tensor
  from collections import OrderedDict
  
  
  logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
                      datefmt='%m/%d/%Y %H:%M:%S',
                      level=logging.INFO)
  logger = logging.getLogger()
  
  
  class CorefModel(nn.Module):
      def __init__(self, config, device, num_genres=None):
          super().__init__()
          self.config = config
          self.eval_only = False
          self.device = device
  
          self.num_genres = num_genres if num_genres else len(config['genres'])
          self.max_seg_len = config['max_segment_len']
          self.max_span_width = config['max_span_width']
          assert config['loss_type'] in ['marginalized', 'hinge']
  
          # Model
          self.dropout = nn.Dropout(p=config['dropout_rate'])
          self.bert = AutoModel.from_pretrained(config['bert_pretrained_name_or_path'])
  
          self.bert_emb_size = self.bert.config.hidden_size
          self.span_emb_size = self.bert_emb_size * 3
          if config['use_features']:
              self.span_emb_size += config['feature_emb_size']
          self.pair_emb_size = self.span_emb_size * 3
          if config['use_metadata']:
              self.pair_emb_size += 2 * config['feature_emb_size']
          if config['use_features']:
              self.pair_emb_size += config['feature_emb_size']
          if config['use_segment_distance']:
              self.pair_emb_size += config['feature_emb_size']
  
          assert config["span_width_embedding_size"] >= self.max_span_width
          self.emb_span_width = self.make_embedding(config["span_width_embedding_size"]) if config['use_features'] else None
          self.emb_span_width_prior = self.make_embedding(config["span_width_embedding_size"]) if config['use_width_prior'] else None
          self.emb_antecedent_distance_prior = self.make_embedding(config['num_antecedent_distance_buckets']) if config['use_distance_prior'] else None
          self.emb_genre = self.make_embedding(self.num_genres)
          self.emb_same_speaker = self.make_embedding(2) if config['use_metadata'] else None
          self.emb_segment_distance = self.make_embedding(config['max_training_sentences']) if config['use_segment_distance'] else None
          self.emb_top_antecedent_distance = self.make_embedding(config['num_antecedent_distance_buckets'])
          self.emb_cluster_size = self.make_embedding(10) if config['fine_grained'] and config['higher_order'] == 'cluster_merging' else None
  
          self.mention_token_attn = self.make_ffnn(self.bert_emb_size, 0, output_size=1) if config['model_heads'] else None
          self.span_emb_score_ffnn = self.make_ffnn(self.span_emb_size, [config['ffnn_size']] * config['ffnn_depth'], output_size=1)
          self.span_width_score_ffnn = self.make_ffnn(config['feature_emb_size'], [config['ffnn_size']] * config['ffnn_depth'], output_size=1) if config['use_width_prior'] else None
          self.coarse_bilinear = self.make_ffnn(self.span_emb_size, 0, output_size=self.span_emb_size)
          self.antecedent_distance_score_ffnn = self.make_ffnn(config['feature_emb_size'], 0, output_size=1) if config['use_distance_prior'] else None
          self.coref_score_ffnn = self.make_ffnn(self.pair_emb_size, [config['ffnn_size']] * config['ffnn_depth'], output_size=1) if config['fine_grained'] else None
  
          self.gate_ffnn = self.make_ffnn(2 * self.span_emb_size, 0, output_size=self.span_emb_size) if config['fine_grained'] and config['coref_depth'] > 1 else None
          self.span_attn_ffnn = self.make_ffnn(self.span_emb_size, 0, output_size=1) if config['fine_grained'] and config['higher_order'] == 'span_clustering' else None
          self.cluster_score_ffnn = self.make_ffnn(3 * self.span_emb_size + config['feature_emb_size'], [config['cluster_ffnn_size']] * config['ffnn_depth'], output_size=1) if config['fine_grained'] and config['higher_order'] == 'cluster_merging' else None
  
          self.update_steps = 0  # Internal use for debug
          self.debug = False
  
      def load_state_dict(self, state_dict: 'OrderedDict[str, Tensor]',
                          strict: bool = True):
          new_shape = (self.config['num_antecedent_distance_buckets'], self.config['feature_emb_size'])
          if 'emb_antecedent_distance_prior.weight' in state_dict and state_dict['emb_antecedent_distance_prior.weight'].shape != new_shape:
              logger.warn('Saved embedding for distance is of different size, some values are initialized randomly.')
              for weight_name in ['emb_antecedent_distance_prior.weight', 'emb_top_antecedent_distance.weight']:
                  state_dict[weight_name] = self.initialize_differently_sized_embedding_layer(
                      state_dict[weight_name],
                      new_shape[0],
                  )
          new_shape = (self.config['max_training_sentences'], self.config['feature_emb_size'])
          emb_segment_distance_weight = 'emb_segment_distance.weight'
          if emb_segment_distance_weight in state_dict and state_dict[emb_segment_distance_weight].shape != new_shape:
              logger.warn('Saved embedding for segment distance is of different size, some values are initialized randomly.')
              state_dict[emb_segment_distance_weight] = self.initialize_differently_sized_embedding_layer(
                  state_dict[emb_segment_distance_weight],
                  new_shape[0],
              )
          return super().load_state_dict(state_dict, strict=strict)
  
      def initialize_differently_sized_embedding_layer(self, old_tensor, new_input_size, std=0.02):
          old_shape = old_tensor.shape
          new_weight = torch.empty(
              (new_input_size, self.config["feature_emb_size"]),
              dtype=old_tensor.dtype
          )
          init.normal_(new_weight, std=std)
          new_weight[
              :min(old_shape[0], new_weight.shape[0]), :min(self.config["feature_emb_size"], new_weight.shape[1])
          ] = old_tensor[:new_weight.shape[0], :new_weight.shape[1]]
          return new_weight
  
      def make_embedding(self, dict_size, std=0.02):
          emb = nn.Embedding(dict_size, self.config['feature_emb_size'])
          init.normal_(emb.weight, std=std)
          return emb
  
      def make_linear(self, in_features, out_features, bias=True, std=0.02):
          linear = nn.Linear(in_features, out_features, bias)
          init.normal_(linear.weight, std=std)
          if bias:
              init.zeros_(linear.bias)
          return linear
  
      def make_ffnn(self, feat_size, hidden_size, output_size):
          if hidden_size is None or hidden_size == 0 or hidden_size == [] or hidden_size == [0]:
              return self.make_linear(feat_size, output_size)
  
          if not isinstance(hidden_size, Iterable):
              hidden_size = [hidden_size]
          ffnn = [self.make_linear(feat_size, hidden_size[0]), nn.ReLU(), self.dropout]
          for i in range(1, len(hidden_size)):
              ffnn += [self.make_linear(hidden_size[i-1], hidden_size[i]), nn.ReLU(), self.dropout]
          ffnn.append(self.make_linear(hidden_size[-1], output_size))
          return nn.Sequential(*ffnn)
  
      def get_params(self, named=False):
          bert_based_param, task_param, incremental_param = [], [], []
          for name, param in self.named_parameters():
              if name.startswith('bert'):
                  to_add = (name, param) if named else param
                  bert_based_param.append(to_add)
              elif 'entity_representation_gate' in name:
                  to_add = (name, param) if named else param
                  incremental_param.append(to_add)
              else:
                  to_add = (name, param) if named else param
                  task_param.append(to_add)
          return bert_based_param, task_param, incremental_param
  
      def get_candidate_spans(self, num_words, sentence_map, gold_info, device='cpu'):
          sentence_indices = sentence_map  # [num tokens]
          candidate_starts = torch.unsqueeze(torch.arange(0, num_words, device=device), 1).repeat(1, self.max_span_width)
          candidate_ends = candidate_starts + torch.arange(0, self.max_span_width, device=device)
          candidate_start_sent_idx = sentence_indices[candidate_starts]
          candidate_end_sent_idx = sentence_indices[torch.min(candidate_ends, torch.tensor(num_words - 1, device=device))]
          candidate_mask = (candidate_ends < num_words) & (candidate_start_sent_idx == candidate_end_sent_idx)
          candidate_starts, candidate_ends = candidate_starts[candidate_mask], candidate_ends[candidate_mask]  # [num valid candidates]
          num_candidates = candidate_starts.shape[0]
  
          candidate_labels = None
          if gold_info is not None:
              same_start = (torch.unsqueeze(gold_info['gold_starts'], 1) == torch.unsqueeze(candidate_starts, 0))
              same_end = (torch.unsqueeze(gold_info['gold_ends'], 1) == torch.unsqueeze(candidate_ends, 0))
              same_span = (same_start & same_end).to(torch.long)
              candidate_labels = torch.matmul(torch.unsqueeze(gold_info['gold_mention_cluster_map'], 0).to(torch.float), same_span.to(torch.float))
              candidate_labels = torch.squeeze(candidate_labels.to(torch.long), 0)  # [num candidates]; non-gold span has label 0
          return {
              'candidate_starts': candidate_starts,
              'candidate_ends': candidate_ends,
              'num_candidates': num_candidates,
              'candidate_labels': candidate_labels,
          }
  
      def get_top_spans(self, candidate_span_emb, candidates, num_words, do_loss, device='cpu'):
          conf = self.config
  
          candidate_starts = candidates['candidate_starts']
          candidate_ends = candidates['candidate_ends']
          candidate_labels = candidates['candidate_labels']
          candidate_width_idx = candidate_ends - candidate_starts
  
          # Get span score
          candidate_mention_scores = torch.squeeze(self.span_emb_score_ffnn(candidate_span_emb), 1)
          if conf['use_width_prior']:
              width_score = torch.squeeze(self.span_width_score_ffnn(self.emb_span_width_prior.weight), 1)
              candidate_width_score = width_score[candidate_width_idx]
              candidate_mention_scores += candidate_width_score
  
          # Extract top spans
          candidate_idx_sorted_by_score = torch.argsort(candidate_mention_scores, descending=True).tolist()
          candidate_starts_cpu, candidate_ends_cpu = candidate_starts.tolist(), candidate_ends.tolist()
          num_top_spans = int(min(conf['max_num_extracted_spans'], conf['top_span_ratio'] * num_words))
          selected_idx_cpu = self._extract_top_spans(candidate_idx_sorted_by_score, candidate_starts_cpu, candidate_ends_cpu, num_top_spans)
          assert len(selected_idx_cpu) == num_top_spans
          selected_idx = torch.tensor(selected_idx_cpu, device=device)
          top_span_starts, top_span_ends = candidate_starts[selected_idx], candidate_ends[selected_idx]
          top_span_emb = candidate_span_emb[selected_idx]
          top_span_cluster_ids = candidate_labels[selected_idx] if do_loss else None
          top_span_mention_scores = candidate_mention_scores[selected_idx]
          return {
              'starts': top_span_starts,
              'ends': top_span_ends,
              'emb': top_span_emb,
              'cluster_ids': top_span_cluster_ids,
              'mention_scores': top_span_mention_scores,
              'candidate_mention_scores': candidate_mention_scores,
              'num_top_spans': num_top_spans,
          }
  
      def build_candidate_embeddings(self, mention_doc, candidate_spans, device='cpu'):
          candidate_ends = candidate_spans['candidate_ends']
          candidate_starts = candidate_spans['candidate_starts']
          num_candidates = candidate_spans['num_candidates']
          num_words = mention_doc.shape[0]
  
          # Get span embedding (these are the bert embeddings)
          span_start_emb, span_end_emb = mention_doc[candidate_starts], mention_doc[candidate_ends]
          # list of all embeddings relevant to the candidate span, they are concated to form the span embedding
          candidate_emb_list = [span_start_emb, span_end_emb]
          if self.config['use_features']:
              candidate_width_idx = candidate_ends - candidate_starts
              candidate_width_emb = self.emb_span_width(candidate_width_idx)
              candidate_width_emb = self.dropout(candidate_width_emb)
              candidate_emb_list.append(candidate_width_emb)
          # Use attended head or avg token
          candidate_tokens = torch.unsqueeze(torch.arange(0, num_words, device=device), 0).repeat(num_candidates, 1)  # [num_candidates, num_words]
          candidate_tokens_mask = (candidate_tokens >= torch.unsqueeze(candidate_starts, 1)) & (candidate_tokens <= torch.unsqueeze(candidate_ends, 1))
          if self.config['model_heads']:
              token_attn = torch.squeeze(self.mention_token_attn(mention_doc), 1)
          else:
              token_attn = torch.ones(num_words, dtype=torch.float, device=device)  # Use avg if no attention
          # Attention is zeroed out where the candidate tokens are False, i.e. where the token is not part of the candidate span
          candidate_tokens_attn_raw = torch.log(candidate_tokens_mask.to(torch.float)) + torch.unsqueeze(token_attn, 0)
          candidate_tokens_attn = nn.functional.softmax(candidate_tokens_attn_raw, dim=1)
          # Attention for each candidate, all tokens in the span are weighted using attention (or if model_heads is False unweighted)
          head_attn_emb = torch.matmul(candidate_tokens_attn, mention_doc)
          candidate_emb_list.append(head_attn_emb)
          return candidate_emb_list
  
      def forward(self, *input):
          return self.get_predictions_and_loss(*input)
  
      def get_predictions_and_loss(self, input_ids, input_mask, speaker_ids, sentence_len, genre, sentence_map,
                                   is_training, gold_starts=None, gold_ends=None, gold_mention_cluster_map=None):
          """ Model and input are already on the device """
          device = self.device
          conf = self.config
  
          do_loss = False
          if gold_mention_cluster_map is not None:
              assert gold_starts is not None
              assert gold_ends is not None
              do_loss = True
  
          # Get token emb
          mention_doc = self.bert(input_ids, attention_mask=input_mask)[0]  # [num seg, num max tokens, emb size]
          input_mask = input_mask.to(torch.bool)
          mention_doc = mention_doc[input_mask]
          speaker_ids = speaker_ids[input_mask]
          num_words = mention_doc.shape[0]
  
          if not self.eval_only and not self.training and sentence_len.shape[0] > conf['doc_max_segments']:
              logger.warn('Not predicting document longer than doc_max_segments')
              return [
                  None,
                  None,
                  None,
                  torch.tensor([]),
                  torch.tensor([]),
                  torch.tensor([]),
                  torch.tensor([[0.1]]),
              ]
  
          # Get candidate span
          if do_loss:
              gold_info = {
                  'gold_starts': gold_starts,
                  'gold_ends': gold_ends,
                  'gold_mention_cluster_map': gold_mention_cluster_map,
              }
          else:
              gold_info = None
          candidate_spans = self.get_candidate_spans(num_words, sentence_map, gold_info, device=device)
          # num_candidates = candidate_spans['num_candidates']
  
          candidate_span_emb = torch.cat(self.build_candidate_embeddings(mention_doc, candidate_spans, device=device), dim=1)  # [num candidates, new emb size]
  
          top_spans = self.get_top_spans(candidate_span_emb, candidate_spans, num_words, do_loss=do_loss, device=device)
          top_span_starts, top_span_ends = top_spans['starts'], top_spans['ends']
          top_span_emb = top_spans['emb']
          top_span_cluster_ids = top_spans['cluster_ids']
          top_span_mention_scores = top_spans['mention_scores']
          num_top_spans = top_spans['num_top_spans']
  
          # Coarse pruning on each mention's antecedents
          max_top_antecedents = min(num_top_spans, conf['max_top_antecedents'])
          top_span_range = torch.arange(0, num_top_spans, device=device)
          antecedent_offsets = torch.unsqueeze(top_span_range, 1) - torch.unsqueeze(top_span_range, 0)
          antecedent_mask = (antecedent_offsets >= 1)
          pairwise_mention_score_sum = torch.unsqueeze(top_span_mention_scores, 1) + torch.unsqueeze(top_span_mention_scores, 0)
          source_span_emb = self.dropout(self.coarse_bilinear(top_span_emb))
          target_span_emb = self.dropout(torch.transpose(top_span_emb, 0, 1))
          pairwise_coref_scores = torch.matmul(source_span_emb, target_span_emb)
          pairwise_fast_scores = pairwise_mention_score_sum + pairwise_coref_scores
          pairwise_fast_scores += torch.log(antecedent_mask.to(torch.float))
          if conf['use_distance_prior']:
              distance_score = torch.squeeze(self.antecedent_distance_score_ffnn(self.dropout(self.emb_antecedent_distance_prior.weight)), 1)
              bucketed_distance = util.bucket_distance(antecedent_offsets, num_buckets=self.config['num_antecedent_distance_buckets'])
              antecedent_distance_score = distance_score[bucketed_distance]
              pairwise_fast_scores += antecedent_distance_score
          top_pairwise_fast_scores, top_antecedent_idx = torch.topk(pairwise_fast_scores, k=max_top_antecedents)
          top_antecedent_mask = util.batch_select(antecedent_mask, top_antecedent_idx, device)  # [num top spans, max top antecedents]
          top_antecedent_offsets = util.batch_select(antecedent_offsets, top_antecedent_idx, device)
  
          # Slow mention ranking
          if conf['fine_grained']:
              same_speaker_emb, genre_emb, seg_distance_emb, top_antecedent_distance_emb = None, None, None, None
              if conf['use_metadata']:
                  top_span_speaker_ids = speaker_ids[top_span_starts]
                  top_antecedent_speaker_id = top_span_speaker_ids[top_antecedent_idx]
                  same_speaker = torch.unsqueeze(top_span_speaker_ids, 1) == top_antecedent_speaker_id
                  same_speaker_emb = self.emb_same_speaker(same_speaker.to(torch.long))
                  genre_emb = self.emb_genre(genre)
                  genre_emb = torch.unsqueeze(torch.unsqueeze(genre_emb, 0), 0).repeat(num_top_spans, max_top_antecedents, 1)
              if conf['use_segment_distance']:
                  num_segs, seg_len = input_ids.shape[0], input_ids.shape[1]
                  token_seg_ids = torch.arange(0, num_segs, device=device).unsqueeze(1).repeat(1, seg_len)
                  token_seg_ids = token_seg_ids[input_mask]
                  top_span_seg_ids = token_seg_ids[top_span_starts]
                  top_antecedent_seg_ids = token_seg_ids[top_span_starts[top_antecedent_idx]]
                  top_antecedent_seg_distance = torch.unsqueeze(top_span_seg_ids, 1) - top_antecedent_seg_ids
                  top_antecedent_seg_distance = torch.clamp(top_antecedent_seg_distance, 0, self.config['max_training_sentences'] - 1)
                  seg_distance_emb = self.emb_segment_distance(top_antecedent_seg_distance)
              if conf['use_features']:  # Antecedent distance
                  top_antecedent_distance = util.bucket_distance(top_antecedent_offsets, num_buckets=self.config['num_antecedent_distance_buckets'])
                  top_antecedent_distance_emb = self.emb_top_antecedent_distance(top_antecedent_distance)
  
              for depth in range(conf['coref_depth']):
                  top_antecedent_emb = top_span_emb[top_antecedent_idx]  # [num top spans, max top antecedents, emb size]
                  feature_list = []
                  if conf['use_metadata']:  # speaker, genre
                      feature_list.append(same_speaker_emb)
                      feature_list.append(genre_emb)
                  if conf['use_segment_distance']:
                      feature_list.append(seg_distance_emb)
                  if conf['use_features']:  # Antecedent distance
                      feature_list.append(top_antecedent_distance_emb)
                  target_emb = torch.unsqueeze(top_span_emb, 1).repeat(1, max_top_antecedents, 1)
                  similarity_emb = target_emb * top_antecedent_emb
                  if len(feature_list) > 0:
                      feature_emb = torch.cat(feature_list, dim=2)
                      feature_emb = self.dropout(feature_emb)
                      pair_emb = torch.cat([target_emb, top_antecedent_emb, similarity_emb, feature_emb], 2)
                  else:
                      pair_emb = torch.cat([target_emb, top_antecedent_emb, similarity_emb], 2)
                  top_pairwise_slow_scores = torch.squeeze(self.coref_score_ffnn(pair_emb), 2)
                  top_pairwise_scores = top_pairwise_slow_scores + top_pairwise_fast_scores
                  if conf['higher_order'] == 'cluster_merging':
                      cluster_merging_scores = ho.cluster_merging(top_span_emb, top_antecedent_idx, top_pairwise_scores, self.emb_cluster_size, self.cluster_score_ffnn, None, self.dropout,
                                                                  device=device, reduce=conf['cluster_reduce'], easy_cluster_first=conf['easy_cluster_first'])
                      break
                  elif depth != conf['coref_depth'] - 1:
                      if conf['higher_order'] == 'attended_antecedent':
                          refined_span_emb = ho.attended_antecedent(top_span_emb, top_antecedent_emb, top_pairwise_scores, device)
                      elif conf['higher_order'] == 'max_antecedent':
                          refined_span_emb = ho.max_antecedent(top_span_emb, top_antecedent_emb, top_pairwise_scores, device)
                      elif conf['higher_order'] == 'entity_equalization':
                          refined_span_emb = ho.entity_equalization(top_span_emb, top_antecedent_emb, top_antecedent_idx, top_pairwise_scores, device)
                      elif conf['higher_order'] == 'span_clustering':
                          refined_span_emb = ho.span_clustering(top_span_emb, top_antecedent_idx, top_pairwise_scores, self.span_attn_ffnn, device)
  
                      gate = self.gate_ffnn(torch.cat([top_span_emb, refined_span_emb], dim=1))
                      gate = torch.sigmoid(gate)
                      top_span_emb = gate * refined_span_emb + (1 - gate) * top_span_emb  # [num top spans, span emb size]
          else:
              top_pairwise_scores = top_pairwise_fast_scores  # [num top spans, max top antecedents]
  
          if not do_loss:
              if conf['fine_grained'] and conf['higher_order'] == 'cluster_merging':
                  top_pairwise_scores += cluster_merging_scores
              top_antecedent_scores = torch.cat([torch.zeros(num_top_spans, 1, device=device), top_pairwise_scores], dim=1)  # [num top spans, max top antecedents + 1]
              return [
                  candidate_spans["candidate_starts"],
                  candidate_spans["candidate_ends"],
                  top_spans["mention_scores"],
                  top_spans["starts"],
                  top_spans["ends"],
                  top_antecedent_idx,
                  top_antecedent_scores,
              ]
  
          # Get gold labels
          top_antecedent_cluster_ids = top_span_cluster_ids[top_antecedent_idx]
          top_antecedent_cluster_ids += (top_antecedent_mask.to(torch.long) - 1) * 100000  # Mask id on invalid antecedents
          same_gold_cluster_indicator = (top_antecedent_cluster_ids == torch.unsqueeze(top_span_cluster_ids, 1))
          non_dummy_indicator = torch.unsqueeze(top_span_cluster_ids > 0, 1)
          pairwise_labels = same_gold_cluster_indicator & non_dummy_indicator
          dummy_antecedent_labels = torch.logical_not(pairwise_labels.any(dim=1, keepdims=True))
          top_antecedent_gold_labels = torch.cat([dummy_antecedent_labels, pairwise_labels], dim=1)
  
          # Get loss
          top_antecedent_scores = torch.cat([torch.zeros(num_top_spans, 1, device=device), top_pairwise_scores], dim=1)
          if conf['loss_type'] == 'marginalized':
              log_marginalized_antecedent_scores = torch.logsumexp(top_antecedent_scores + torch.log(top_antecedent_gold_labels.to(torch.float)), dim=1)
              log_norm = torch.logsumexp(top_antecedent_scores, dim=1)
              loss = torch.sum(log_norm - log_marginalized_antecedent_scores)
          elif conf['loss_type'] == 'hinge':
              top_antecedent_mask = torch.cat([torch.ones(num_top_spans, 1, dtype=torch.bool, device=device), top_antecedent_mask], dim=1)
              top_antecedent_scores += torch.log(top_antecedent_mask.to(torch.float))
              highest_antecedent_scores, highest_antecedent_idx = torch.max(top_antecedent_scores, dim=1)
              gold_antecedent_scores = top_antecedent_scores + torch.log(top_antecedent_gold_labels.to(torch.float))
              highest_gold_antecedent_scores, highest_gold_antecedent_idx = torch.max(gold_antecedent_scores, dim=1)
              slack_hinge = 1 + highest_antecedent_scores - highest_gold_antecedent_scores
              # Calculate delta
              highest_antecedent_is_gold = (highest_antecedent_idx == highest_gold_antecedent_idx)
              mistake_false_new = (highest_antecedent_idx == 0) & torch.logical_not(dummy_antecedent_labels.squeeze())
              delta = ((3 - conf['false_new_delta']) / 2) * torch.ones(num_top_spans, dtype=torch.float, device=device)
              delta -= (1 - conf['false_new_delta']) * mistake_false_new.to(torch.float)
              delta *= torch.logical_not(highest_antecedent_is_gold).to(torch.float)
              loss = torch.sum(slack_hinge * delta)
  
          # Add mention loss
          if conf['mention_loss_coef']:
              gold_mention_scores = top_span_mention_scores[top_span_cluster_ids > 0]
              loss_mention = -torch.sum(torch.log(torch.sigmoid(gold_mention_scores))) * conf['mention_loss_coef']
              loss += loss_mention
  
          if conf['higher_order'] == 'cluster_merging':
              top_pairwise_scores += cluster_merging_scores
              top_antecedent_scores = torch.cat([torch.zeros(num_top_spans, 1, device=device), top_pairwise_scores], dim=1)
              log_marginalized_antecedent_scores2 = torch.logsumexp(top_antecedent_scores + torch.log(top_antecedent_gold_labels.to(torch.float)), dim=1)
              log_norm2 = torch.logsumexp(top_antecedent_scores, dim=1)  # [num top spans]
              loss_cm = torch.sum(log_norm2 - log_marginalized_antecedent_scores2)
              if conf['cluster_dloss']:
                  loss += loss_cm
              else:
                  loss = loss_cm
  
          # Debug
          if self.debug:
              if self.update_steps % 20 == 0:
                  logger.info('---------debug step: %d---------' % self.update_steps)
                  # logger.info('candidates: %d; antecedents: %d' % (num_candidates, max_top_antecedents))
                  logger.info('spans/gold: %d/%d; ratio: %.2f' % (num_top_spans, (top_span_cluster_ids > 0).cpu().numpy().sum(), (top_span_cluster_ids > 0).cpu().numpy().sum()/num_top_spans))
                  if conf['mention_loss_coef']:
                      logger.info('mention loss: %.4f' % loss_mention.item())
                  if conf['loss_type'] == 'marginalized':
                      logger.info('norm/gold: %.4f/%.4f' % (torch.sum(log_norm).item(), torch.sum(log_marginalized_antecedent_scores).item()))
                  else:
                      logger.info('loss: %.4f' % loss.item())
          self.update_steps += 1
  
          return [
              candidate_spans['candidate_starts'],
              candidate_spans['candidate_ends'],
              top_spans['candidate_mention_scores'],
              top_spans['starts'],
              top_spans['ends'],
              top_antecedent_idx,
              top_antecedent_scores
          ], loss
  
      def _extract_top_spans(self, candidate_idx_sorted, candidate_starts, candidate_ends, num_top_spans):
          """ Keep top non-cross-overlapping candidates ordered by scores; compute on CPU because of loop """
          selected_candidate_idx = []
          start_to_max_end, end_to_min_start = {}, {}
          for candidate_idx in candidate_idx_sorted:
              if len(selected_candidate_idx) >= num_top_spans:
                  break
              # Perform overlapping check
              span_start_idx = candidate_starts[candidate_idx]
              span_end_idx = candidate_ends[candidate_idx]
              cross_overlap = False
              for token_idx in range(span_start_idx, span_end_idx + 1):
                  max_end = start_to_max_end.get(token_idx, -1)
                  if token_idx > span_start_idx and max_end > span_end_idx:
                      cross_overlap = True
                      break
                  min_start = end_to_min_start.get(token_idx, -1)
                  if token_idx < span_end_idx and 0 <= min_start < span_start_idx:
                      cross_overlap = True
                      break
              if not cross_overlap:
                  # Pass check; select idx and update dict stats
                  selected_candidate_idx.append(candidate_idx)
                  max_end = start_to_max_end.get(span_start_idx, -1)
                  if span_end_idx > max_end:
                      start_to_max_end[span_start_idx] = span_end_idx
                  min_start = end_to_min_start.get(span_end_idx, -1)
                  if min_start == -1 or span_start_idx < min_start:
                      end_to_min_start[span_end_idx] = span_start_idx
          # Sort selected candidates by span idx
          selected_candidate_idx = sorted(selected_candidate_idx, key=lambda idx: (candidate_starts[idx], candidate_ends[idx]))
          if len(selected_candidate_idx) < num_top_spans:  # Padding
              selected_candidate_idx += ([selected_candidate_idx[0]] * (num_top_spans - len(selected_candidate_idx)))
          return selected_candidate_idx
  
      def get_predicted_antecedents(self, antecedent_idx, antecedent_scores):
          """ CPU list input """
          predicted_antecedents = []
          for i, idx in enumerate(np.argmax(antecedent_scores, axis=1) - 1):
              if idx < 0:
                  predicted_antecedents.append(-1)
              else:
                  predicted_antecedents.append(antecedent_idx[i][idx])
          return predicted_antecedents
  
      def get_predicted_clusters(self, span_starts, span_ends, antecedent_idx, antecedent_scores):
          """ CPU list input """
          # Get predicted antecedents
          predicted_antecedents = self.get_predicted_antecedents(antecedent_idx, antecedent_scores)
  
          # Get predicted clusters
          mention_to_cluster_id = {}
          predicted_clusters = []
          for i, predicted_idx in enumerate(predicted_antecedents):
              if predicted_idx < 0:
                  continue
              assert i > predicted_idx, f'span idx: {i}; antecedent idx: {predicted_idx}'
              # Check antecedent's cluster
              antecedent = (int(span_starts[predicted_idx]), int(span_ends[predicted_idx]))
              antecedent_cluster_id = mention_to_cluster_id.get(antecedent, -1)
              if antecedent_cluster_id == -1:
                  antecedent_cluster_id = len(predicted_clusters)
                  predicted_clusters.append([antecedent])
                  mention_to_cluster_id[antecedent] = antecedent_cluster_id
              # Add mention to cluster
              mention = (int(span_starts[i]), int(span_ends[i]))
              predicted_clusters[antecedent_cluster_id].append(mention)
              mention_to_cluster_id[mention] = antecedent_cluster_id
  
          predicted_clusters = [tuple(c) for c in predicted_clusters]
          return predicted_clusters, mention_to_cluster_id, predicted_antecedents
  
      def update_evaluator(self, span_starts, span_ends, antecedent_idx, antecedent_scores, gold_clusters, evaluator):
          predicted_clusters, mention_to_cluster_id, _ = self.get_predicted_clusters(span_starts, span_ends, antecedent_idx, antecedent_scores)
          mention_to_predicted = {m: predicted_clusters[cluster_idx] for m, cluster_idx in mention_to_cluster_id.items()}
          gold_clusters = [tuple(tuple(m) for m in cluster) for cluster in gold_clusters]
          mention_to_gold = {m: cluster for cluster in gold_clusters for m in cluster}
          evaluator.update(predicted_clusters, gold_clusters, mention_to_predicted, mention_to_gold)
          return predicted_clusters
  
  
  class MentionModel(CorefModel):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.loss = torch.nn.BCEWithLogitsLoss(torch.tensor(self.config["positive_class_weight"]))
  
      def forward(self, input_ids, input_mask, speaker_ids, sentence_len, genre, sentence_map,
                  is_training, gold_starts=None, gold_ends=None, gold_mention_cluster_map=None):
          device = self.device
          conf = self.config
  
          if sentence_len.shape[0] > conf['doc_max_segments']:
              logger.warn('Not predicting document longer than doc_max_segments')
              return [
                  torch.tensor([], device=device),
                  torch.tensor([], device=device),
                  {},
                  torch.tensor(0.0, requires_grad=True),
              ]
  
          mention_doc = self.bert(input_ids, attention_mask=input_mask)[0]  # [num seg, num max tokens, emb size]
  
          input_mask = input_mask.to(torch.bool)
          mention_doc = mention_doc[input_mask]
          speaker_ids = speaker_ids[input_mask]
  
          num_words = mention_doc.shape[0]
          gold_info = {
              "gold_starts": gold_starts,
              "gold_ends": gold_ends,
              "gold_mention_cluster_map": gold_mention_cluster_map
          }
          candidates = self.get_candidate_spans(num_words, sentence_map, gold_info, device=device)
  
          span_start_emb, span_end_emb = mention_doc[candidates['candidate_starts']], mention_doc[candidates['candidate_ends']]
          candidate_emb_list = [span_start_emb, span_end_emb]
  
          if conf['use_features']:
              candidate_width_idx = candidates['candidate_ends'] - candidates['candidate_starts']
              candidate_width_emb = self.emb_span_width(candidate_width_idx)
              candidate_width_emb = self.dropout(candidate_width_emb)
              candidate_emb_list.append(candidate_width_emb)
          candidate_tokens = torch.unsqueeze(torch.arange(0, num_words, device=device), 0).repeat(candidates['num_candidates'], 1)
          candidate_tokens_mask = (candidate_tokens >= torch.unsqueeze(candidates['candidate_starts'], 1)) & \
              (candidate_tokens <= torch.unsqueeze(candidates['candidate_ends'], 1))
          if conf['model_heads']:
              token_attn = torch.squeeze(self.mention_token_attn(mention_doc), 1)
          else:
              token_attn = torch.ones(num_words, dtype=torch.float, device=device)  # Use avg if no attention
          candidate_tokens_attn_raw = torch.log(candidate_tokens_mask.to(torch.float)) + torch.unsqueeze(token_attn, 0)
          candidate_tokens_attn = nn.functional.softmax(candidate_tokens_attn_raw, dim=1)
          head_attn_emb = torch.matmul(candidate_tokens_attn, mention_doc)
          candidate_emb_list.append(head_attn_emb)
          candidate_span_emb = torch.cat(candidate_emb_list, dim=1)  # [num candidates, new emb size]
          scores = self.span_emb_score_ffnn(candidate_span_emb)
          loss = self.loss(scores.squeeze(), candidates["candidate_labels"].clone().to(torch.bool).to(torch.float))
          cluster_to_spans = defaultdict(list)
          new_cluster = 0
          for label, score, start, end in zip(
                  candidates["candidate_labels"],
                  scores,
                  candidates["candidate_starts"],
                  candidates["candidate_ends"],
          ):
              if score > 0.5:
                  new_cluster += 1
                  cluster_to_spans[new_cluster].append((start.item(), end.item()))
  
          return scores, candidates["candidate_labels"], cluster_to_spans, loss
  
  
  class IncrementalCorefModel(CorefModel):
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.loss = torch.nn.CrossEntropyLoss()
          conf = args[0]
          # This is a really good idea in incremental models, uncomment at your peril
          # Otherwise only sub-sections of your documents are used, not great...
          assert conf['long_doc_strategy'] == 'keep'
  
          # Takes concat(entity_representation, span_representation)
          self.entity_representation_gate = nn.Linear(self.span_emb_size * 2, 1)
          # self.create_entity = torch.nn.Embedding(1, self.config['feature_emb_size'])
  
      def forward(self, *input, **kwargs):
          return self.get_predictions_incremental(*input, **kwargs)
  
      def update_evaluator(self, span_starts, span_ends, predicted_clusters, mention_to_cluster_id, gold_clusters, evaluator):
          mention_to_predicted = {m: tuple(predicted_clusters[cluster_idx]) for m, cluster_idx in mention_to_cluster_id.items()}
          gold_clusters = [tuple(tuple(m) for m in cluster) for cluster in gold_clusters]
          mention_to_gold = {m: cluster for cluster in gold_clusters for m in cluster}
          evaluator.update(predicted_clusters, gold_clusters, mention_to_predicted, mention_to_gold)
  
      def get_predictions_incremental(self, input_ids, input_mask, speaker_ids, sentence_len, genre, sentence_map,
                                      is_training, gold_starts=None, gold_ends=None, gold_mention_cluster_map=None,
                                      global_loss_chance=0.0, teacher_forcing=False):
          max_segments = 5
  
          mention_to_cluster_id = {}
          predicted_clusters = []
          entities = None
          if torch.rand(1) < global_loss_chance:
              loss_strategy = GoldLabelStrategy.ORIGINAL
          else:
              loss_strategy = GoldLabelStrategy.MOST_RECENT
          cpu_entities = IncrementalEntities(conf=self.config, device="cpu", gold_strategy=loss_strategy)
  
          do_loss = False
          if gold_mention_cluster_map is not None:
              assert gold_starts is not None
              assert gold_ends is not None
              do_loss = True
  
          offset = 0
          total_loss = torch.tensor([0.0], requires_grad=True, device=self.device)
          total_gold_mask = None
          for i, start in enumerate(range(0, input_ids.shape[0], max_segments)):
              end = start + max_segments
              start_offset = torch.sum(input_mask[:start], (0, 1))
              delta_offset = torch.sum(input_mask[start:end], (0, 1))
              end_offset = start_offset + delta_offset
              if gold_ends is not None:
                  # We include any gold spans that either end or start in the current window
                  gold_mask = (gold_starts < end_offset) & (gold_starts > start_offset) | ((gold_ends > start_offset) & (gold_ends < end_offset))
                  if total_gold_mask is None:
                      total_gold_mask = gold_mask
                  else:
                      total_gold_mask |= gold_mask
                  windowed_gold_starts = gold_starts[gold_mask].clamp(start_offset, end_offset) - start_offset
                  windowed_gold_ends = gold_ends[gold_mask].clamp(start_offset, end_offset) - start_offset
                  windowed_gold_mention_cluster_map = gold_mention_cluster_map[gold_mask]
              else:
                  windowed_gold_starts = None
                  windowed_gold_ends = None
                  windowed_gold_mention_cluster_map = None
              res = self.get_predictions_incremental_internal(
                  input_ids[start:end],
                  input_mask[start:end],
                  speaker_ids[start:end],
                  sentence_len[start:end],
                  genre,
                  sentence_map[start_offset:end_offset],
                  is_training,
                  gold_starts=windowed_gold_starts,
                  gold_ends=windowed_gold_ends,
                  gold_mention_cluster_map=windowed_gold_mention_cluster_map,
                  entities=entities,
                  do_loss=do_loss,
                  offset=offset,
                  loss_strategy=loss_strategy,
                  teacher_forcing=teacher_forcing,
              )
              offset += torch.sum(input_mask[start:end], (0, 1)).item()
              if do_loss:
                  entities, new_cpu_entities, loss = res
                  total_loss = loss + total_loss
              else:
                  entities, new_cpu_entities = res
              cpu_entities.extend(new_cpu_entities)
          cpu_entities.extend(entities)
          starts, ends, mention_to_cluster_id, predicted_clusters = cpu_entities.get_result(
              remove_singletons=not self.config['incremental_singletons']
          )
          out = [
              starts,
              ends,
              mention_to_cluster_id,
              predicted_clusters,
          ]
          if do_loss:
              return out, loss
          else:
              return out
  
      def get_predictions_incremental_internal(self, input_ids, input_mask, speaker_ids, sentence_len, genre, sentence_map,
                                               is_training, gold_starts=None, gold_ends=None, gold_mention_cluster_map=None,
                                               entities=None, do_loss=None, offset=0, loss_strategy=GoldLabelStrategy.MOST_RECENT,
                                               teacher_forcing=False):
          device = self.device
          conf = self.config
          return_singletons = conf['incremental_singletons']
  
          # The model should already be trained so we detach the BERT part, massively improving performance
          mention_doc = self.bert(input_ids, attention_mask=input_mask)[0].detach()  # [num seg, num max tokens, emb size]
  
          input_mask = input_mask.to(torch.bool)
          mention_doc = mention_doc[input_mask]
          speaker_ids = speaker_ids[input_mask]
  
          num_words = mention_doc.shape[0]
  
          if do_loss:
              gold_info = {
                  'gold_starts': gold_starts,
                  'gold_ends': gold_ends,
                  'gold_mention_cluster_map': gold_mention_cluster_map,
              }
              labels_for_starts = {(s.item(), e.item()): v.item() for s, e, v in zip(gold_starts, gold_ends, gold_mention_cluster_map)}
          else:
              gold_info = None
              labels_for_starts = {}
  
          candidate_spans = self.get_candidate_spans(num_words, sentence_map, gold_info, device=device)
          candidate_span_emb = torch.cat(self.build_candidate_embeddings(mention_doc, candidate_spans, device=device), dim=1)  # [num candidates, new emb size]
  
          top_spans = self.get_top_spans(candidate_span_emb, candidate_spans, num_words, do_loss=do_loss, device=device)
          top_span_starts, top_span_ends = top_spans['starts'], top_spans['ends']
          top_span_emb = top_spans['emb']
  
          new_cluster_threshold = torch.tensor([conf['new_cluster_threshold']]).unsqueeze(0).to(device)
          cpu_entities = IncrementalEntities(conf, "cpu", gold_strategy=loss_strategy)
          if entities is None:
              entities = IncrementalEntities(conf, device, gold_strategy=loss_strategy)
  
          if len(top_span_emb.shape) == 1:
              top_span_emb = top_span_emb.unsqueeze(0)
  
          losses = []
          cpu_loss = 0.0
          discard_weight = len(labels_for_starts.keys()) / len(top_span_starts)
          new_cluster_weight = len(set(labels_for_starts.keys())) / len(top_span_starts)
          for emb, span_start, span_end, mention_score in zip(top_span_emb, top_span_starts, top_span_ends, top_spans['mention_scores']):
              gold_class = labels_for_starts.get((span_start.item(), span_end.item()))
              if len(entities) == 0:
                  # No need to do the whole similarity computation, this is the first mention
                  entities.add_entity(emb, gold_class, span_start, span_end, offset=offset)
              else:
                  if conf['evict']:
                      entities.evict(evict_to=cpu_entities)
                  feature_list = []
                  if conf['use_metadata']:
                      same_speaker_emb = torch.zeros(conf['feature_emb_size'], device=self.device)
                      genre_emb = self.emb_genre(genre)
                      feature_list.append(same_speaker_emb.repeat(entities.emb.shape[0]).reshape(-1, conf['feature_emb_size']))
                      feature_list.append(genre_emb.repeat(feature_list[0].shape[0]).reshape(-1, conf['feature_emb_size']))
                  if conf['use_segment_distance']:
                      seg_distance_emb = self.emb_segment_distance(entities.sentence_distance.type(torch.long))
                      feature_list.append(seg_distance_emb.reshape(-1, conf['feature_emb_size']))
                  if conf['use_features']:
                      dists = util.bucket_distance(entities.mention_distance, num_buckets=self.config['num_antecedent_distance_buckets'])
                      antecedent_distance_emb = self.emb_top_antecedent_distance(dists.type(torch.long))
                      feature_list.append(antecedent_distance_emb.reshape(-1, conf['feature_emb_size']))
                  fast_source_span_emb = self.dropout(self.coarse_bilinear(emb))
                  fast_entity_embs = self.dropout(torch.transpose(entities.emb, 0, 1))
                  fast_coref_scores = torch.matmul(fast_source_span_emb, fast_entity_embs).unsqueeze(-1)
                  feature_emb = torch.cat(feature_list, dim=1)
                  feature_emb = self.dropout(feature_emb)
                  embs = emb.repeat(entities.emb.shape[0], 1)
                  similarity_emb = embs * entities.emb
                  pair_emb = torch.cat([embs, entities.emb, similarity_emb, feature_emb], 1)
                  # It's important for us to also involve the mention span scores, the only way to prune discovered spans is by turning them into singleton clusters
                  # This is encouraged by explicitly involving the sum of the mention scores
                  original_scores = self.coref_score_ffnn(pair_emb) + fast_coref_scores
                  if return_singletons:
                      scores = torch.cat([new_cluster_threshold, -mention_score.view(1, 1), original_scores])
                  else:
                      scores = torch.cat([new_cluster_threshold, original_scores])
                  dist = torch.softmax(scores, 0)
  
                  index_to_update = dist.argmax()
                  if return_singletons:
                      cluster_to_update = index_to_update - 2
                  else:
                      cluster_to_update = index_to_update - 1
                  weights = torch.ones(scores.squeeze().T.shape)
                  if return_singletons:
                      weights[0] = new_cluster_weight
                      weights[1] = discard_weight
                  else:
                      weights[0] = discard_weight
                  cre_loss = torch.nn.CrossEntropyLoss(weight=weights.to(self.device))
                  if gold_class and do_loss:
                      if return_singletons:
                          target = torch.tensor([entities.class_gold_entity.get(gold_class, -2) + 2]).to(device)
                      else:
                          target = torch.tensor([entities.class_gold_entity.get(gold_class, -1) + 1]).to(device)
                      loss = cre_loss(scores.T, target)
                      losses.append(loss)
                  elif do_loss:
                      # In this case we are training but don't have a gold label for this span
                      # i.e. the span is not in any gold cluster!
                      if return_singletons:
                          # In this case we add the option to discard a mention
                          loss = cre_loss(scores.T, torch.tensor([1], device=device))
                      else:
                          # Always create a new singleton cluster hoping nothing else ever gets added
                          loss = cre_loss(scores.T, torch.tensor([0], device=device))
                      losses.append(loss)
                  if util.cuda_allocated_memory(self.device) > conf['memory_limit'] and len(losses) > 0:
                      sum(losses).backward(retain_graph=True)
                      cpu_loss += sum(losses).item()
                      losses = []
                  if teacher_forcing:
                      forced_class = entities.class_gold_entity.get(gold_class)
                      if forced_class is None:
                          cluster_to_update = None
                      else:
                          cluster_to_update = torch.tensor(forced_class)
                      if cluster_to_update is None:
                          index_to_update = 0
                  if index_to_update == 0:
                      entities.add_entity(emb, gold_class, span_start, span_end, offset=offset)
                  elif index_to_update == 1 and return_singletons:
                      pass
                  else:
                      update_gate = torch.sigmoid(
                          self.entity_representation_gate(torch.cat(
                              [
                                  emb,
                                  entities.emb[cluster_to_update],
                              ],
                          )))
                      entities.update_entity(
                          cluster_to_update,
                          emb,
                          gold_class,
                          span_start,
                          span_end,
                          update_gate,
                          offset=offset
                      )
  
          if len(losses) > 0:
              sum(losses).backward(retain_graph=True)
              cpu_loss += sum(losses).item()
          if do_loss:
              return entities, cpu_entities, cpu_loss
          else:
              return entities, cpu_entities
  