inference_utils.py

# Inference functions for the SRL model.
import numpy as np


# TODO: Pass arg 
def decode_spans(span_starts, span_ends, span_scores, labels_inv):
  """
  Args:
    span_starts: [num_candidates,]
    span_scores: [num_candidates, num_labels]
  """
  pred_spans = []
  span_labels = np.argmax(span_scores, axis=1)  # [num_candidates]
  spans_list = zip(span_starts, span_ends, span_labels, span_scores)
  spans_list = sorted(spans_list, key=lambda x: x[3][x[2]], reverse=True)
  predicted_spans = {}
  for start, end, label, _ in spans_list:
    # Skip invalid span.
    if label == 0 or (start, end) in predicted_spans:
      continue
    pred_spans.append((start, end, labels_inv[label]))
    predicted_spans[(start, end)] = label
  return pred_spans
     

def greedy_decode(predict_dict, srl_labels_inv):
  """Greedy decoding for SRL predicate-argument structures.

  Args:
    predict_dict: Dictionary of name to numpy arrays.
    srl_labels_inv: SRL label id to string name.
    suppress_overlap: Whether to greedily suppress overlapping arguments for the same predicate.

  Returns:
    A dictionary from predicate ids to lists of argument spans.
  """
  arg_starts = predict_dict["arg_starts"]
  arg_ends = predict_dict["arg_ends"]
  predicates = predict_dict["predicates"]
  arg_labels = predict_dict["arg_labels"]
  scores = predict_dict["srl_scores"]

  num_suppressed_args = 0
 
  # Map from predicates to a list of labeled spans.
  pred_to_args = {}
  if len(arg_ends) > 0 and len(predicates) > 0:
    max_len = max(np.max(arg_ends), np.max(predicates)) + 1
  else:
    max_len = 1

  for j, pred_id in enumerate(predicates):
    args_list = []
    for i, (arg_start, arg_end) in enumerate(zip(arg_starts, arg_ends)):
      # If label is not null.
      if arg_labels[i][j] == 0:
        continue
      label = srl_labels_inv[arg_labels[i][j]]
      #if label not in ["V", "C-V"]:
      args_list.append((arg_start, arg_end, label, scores[i][j][arg_labels[i][j]]))

    # Sort arguments by highest score first.
    args_list = sorted(args_list, key=lambda x: x[3], reverse=True)
    new_args_list = []

    flags = [False for _ in range(max_len)]
    # Predicate will not overlap with arguments either. 
    flags[pred_id] = True

    for (arg_start, arg_end, label, score) in args_list:
      # If none of the tokens has been covered:
      if not max(flags[arg_start:arg_end+1]):
        new_args_list.append((arg_start, arg_end, label))
        for k in range(arg_start, arg_end+1):
          flags[k] = True

    # Only add predicate if it has any argument.
    if new_args_list:
      pred_to_args[pred_id] = new_args_list

    num_suppressed_args += len(args_list) - len(new_args_list)

  return pred_to_args, num_suppressed_args 
  

_CORE_ARGS = { "ARG0": 1, "ARG1": 2, "ARG2": 4, "ARG3": 8, "ARG4": 16, "ARG5": 32, "ARGA": 64,
               "A0": 1, "A1": 2, "A2": 4, "A3": 8, "A4": 16, "A5": 32, "AA": 64 }


def dp_decode(predict_dict, srl_labels_inv):
  """Decode arguments with dynamic programming. Enforce two constraints:
  1. Non-overlapping constraint.
  2. Unique core arg constraint.
  """
  arg_starts = predict_dict["arg_starts"]
  arg_ends = predict_dict["arg_ends"]
  predicates = predict_dict["predicates"]
  # Greedy labels.
  arg_labels = predict_dict["arg_labels"]  # [args,predicates]
  scores = predict_dict["srl_scores"]  # [args, predicates, roles] 

  pred_to_args = {}
  num_roles = scores.shape[2]

  for j, pred_id in enumerate(predicates):
    num_args = len(arg_starts)
    args = zip(arg_starts, arg_ends, range(num_args))
    args = sorted(args, key=lambda x: (x[0], x[1]))
    #print args

    text_len = max(max(arg_ends), pred_id) + 2
    f = np.log(np.zeros([text_len, 128], dtype=float))
    f[0,0] = .0
    states = { 0: set([0]) }  # A dictionary from id to list of binary core-arg states.
    pointers = {}  # A dictionary from states to (arg_id, role, prev_t, prev_rs) 
    best_state = [(0, 0)]

    def _update_state(t0, rs0, t1, rs1, delta, arg_id, role):
      if f[t0][rs0] + delta > f[t1][rs1]:
        f[t1][rs1] = f[t0][rs0] + delta
        if t1 not in states:
          states[t1] = set()
        states[t1].update([rs1])
        pointers[(t1, rs1)] = (arg_id, role, t0, rs0)
        if f[t1][rs1] > f[best_state[0][0]][best_state[0][1]]:
          best_state[0] = (t1, rs1)

    for start, end, i in args:
      assert scores[i][j][0] == 0
      # The extra dummy score should be same for all states, so we can safely skip arguments overlap
      # with the predicate.
      if start <= pred_id and pred_id <= end:
        continue
      # Locally best role assignment.
      r0 = arg_labels[i][j]
      # Strictly better to incorporate a dummy span if it has the highest local score.
      if r0 == 0:
        continue
      r0_str = srl_labels_inv[r0]
      # Enumerate explored states.
      t_states = [t for t in states.keys() if t <= start]
      for t in t_states:
        role_states = states[t]
        # Update states if best role is not a core arg.
        if not r0_str in _CORE_ARGS:
          for rs in role_states:
            _update_state(t, rs, end+1, rs, scores[i][j][r0], i, r0)
        else:
          core_state = _CORE_ARGS[r0_str]
          # Get highest-scored non-core arg.
          r1 = 0
          for r in range(1, num_roles):
            if scores[i][j][r] > scores[i][j][r1] and srl_labels_inv[r] not in _CORE_ARGS:
              r1 = r
          for rs in role_states:
            #print i, t, rs, core_state, r0, r0_str, r1
            if core_state & rs == 0:
              _update_state(t, rs, end+1, rs|core_state, scores[i][j][r0], i, r0)
            elif r1 > 0: 
              _update_state(t, rs, end+1, rs, scores[i][j][r0], i, r1)
    '''print f
    print states
    print pointers'''
    # Backtrack to decode.
    args_list = []
    t, rs = best_state[0]
    while (t, rs) in pointers:
      i, r, t0, rs0 = pointers[(t, rs)]
      args_list.append((arg_starts[i], arg_ends[i], srl_labels_inv[r]))
      t = t0
      rs = rs0
    if args_list:
      pred_to_args[pred_id] = args_list[::-1]
       
  return pred_to_args, 0


# Coref decoding.
def get_predicted_antecedents(antecedents, antecedent_scores):
  predicted_antecedents = []
  for i, index in enumerate(np.argmax(antecedent_scores, axis=1) - 1):
    if index < 0:
      predicted_antecedents.append(-1)
    else:
      predicted_antecedents.append(antecedents[i, index])
  return predicted_antecedents


def get_predicted_clusters(top_span_starts, top_span_ends, predicted_antecedents):
  mention_to_predicted = {}
  predicted_clusters = []
  for i, predicted_index in enumerate(predicted_antecedents):
    if predicted_index < 0:
      continue
    assert i > predicted_index
    predicted_antecedent = (int(top_span_starts[predicted_index]), int(top_span_ends[predicted_index]))
    if predicted_antecedent in mention_to_predicted:
      predicted_cluster = mention_to_predicted[predicted_antecedent]
    else:
      predicted_cluster = len(predicted_clusters)
      predicted_clusters.append([predicted_antecedent])
      mention_to_predicted[predicted_antecedent] = predicted_cluster

    mention = (int(top_span_starts[i]), int(top_span_ends[i]))
    predicted_clusters[predicted_cluster].append(mention)
    mention_to_predicted[mention] = predicted_cluster

  predicted_clusters = [tuple(pc) for pc in predicted_clusters]
  mention_to_predicted = { m:predicted_clusters[i] for m,i in mention_to_predicted.items() }

  return predicted_clusters, mention_to_predicted


def _decode_non_overlapping_spans(starts, ends, scores, max_len, labels_inv, pred_id):
  labels = np.argmax(scores, axis=1)
  spans = []
  for i, (start, end, label) in enumerate(zip(starts, ends, labels)):
    if label <= 0:
      continue
    label_str = labels_inv[label]
    if pred_id is not None and label_str == "V":
      continue
    spans.append((start, end, label_str, scores[i][label]))
  spans = sorted(spans, key=lambda x: x[3], reverse=True)
  flags = np.zeros([max_len], dtype=bool)
  if pred_id is not None:
    flags[pred_id] = True
  new_spans = []
  for start, end, label_str, score in spans:
    if not max(flags[start:end+1]):
      new_spans.append((start, end, label_str)) #, score))
      for k in range(start, end+1):
        flags[k] = True
  return new_spans


def _dp_decode_non_overlapping_spans(starts, ends, scores, max_len, labels_inv, pred_id, u_constraint=False):
  num_roles = scores.shape[1]
  labels = np.argmax(scores, axis=1)
  spans = zip(starts, ends, range(len(starts)))
  spans = sorted(spans, key=lambda x: (x[0], x[1]))

  if u_constraint:
    f = np.zeros([max_len + 1, 128], dtype=float) - 0.1
  else:
    f = np.zeros([max_len + 1, 1], dtype=float) - 0.1
  f[0, 0] = 0
  states = { 0: set([0]) }  # A dictionary from id to list of binary core-arg states.
  pointers = {}  # A dictionary from states to (arg_id, role, prev_t, prev_rs) 
  best_state = [(0, 0)]

  def _update_state(t0, rs0, t1, rs1, delta, arg_id, role):
    if f[t0][rs0] + delta > f[t1][rs1]:
      f[t1][rs1] = f[t0][rs0] + delta
      if t1 not in states:
        states[t1] = set()
      states[t1].update([rs1])
      pointers[(t1, rs1)] = (arg_id, role, t0, rs0)
      if f[t1][rs1] > f[best_state[0][0]][best_state[0][1]]:
        best_state[0] = (t1, rs1)

  for start, end, i in spans:
    assert scores[i][0] == 0
    # The extra dummy score should be same for all states, so we can safely skip arguments overlap
    # with the predicate.
    if pred_id is not None and start <= pred_id and pred_id <= end:
      continue
    r0 = labels[i]  # Locally best role assignment.
    # Strictly better to incorporate a dummy span if it has the highest local score.
    if r0 == 0:
      continue
    r0_str = labels_inv[r0]
    # Enumerate explored states.
    t_states = [t for t in states.keys() if t <= start]
    for t in t_states:
      role_states = states[t]
      # Update states if best role is not a core arg.
      if not u_constraint or not r0_str in _CORE_ARGS:
        for rs in role_states:
          _update_state(t, rs, end+1, rs, scores[i][r0], i, r0)
      else:
        for rs in role_states:
          for r in range(1, num_roles):
            if scores[i][r] > 0:
              r_str = labels_inv[r]
              core_state = _CORE_ARGS.get(r_str, 0)
              #print start, end, i, r_str, core_state, rs
              if core_state & rs == 0:            
                _update_state(t, rs, end+1, rs|core_state, scores[i][r], i, r)
  # Backtrack to decode.
  new_spans = []
  t, rs = best_state[0]
  while (t, rs) in pointers:
    i, r, t0, rs0 = pointers[(t, rs)]
    new_spans.append((starts[i], ends[i], labels_inv[r]))
    t = t0
    rs = rs0

  return new_spans[::-1]


def srl_decode(sentences, predict_dict, srl_labels_inv, config):
  predictions = {}
  
  # Decode sentence-level tasks.
  num_sentences = len(sentences)
  predictions["srl"] = [{} for i in range(num_sentences)]
      
  # Sentence-level predictions.
  for i in range(num_sentences):
    if "srl" in predictions:
      num_args = predict_dict["num_args"][i]
      num_preds = predict_dict["num_preds"][i]
      for j, pred_id in enumerate(predict_dict["predicates"][i][:num_preds]):
        #arg_spans = _decode_non_overlapping_spans(
        arg_spans = _dp_decode_non_overlapping_spans(
            predict_dict["arg_starts"][i][:num_args],
            predict_dict["arg_ends"][i][:num_args],
            predict_dict["srl_scores"][i,:num_args,j,:],
            len(sentences[i]), srl_labels_inv, pred_id, config["enforce_srl_constraint"])
        # To avoid warnings in the eval script.
        if config["use_gold_predicates"]:
          arg_spans.append((pred_id, pred_id, "V"))
        if arg_spans:
          predictions["srl"][i][pred_id] = sorted(arg_spans, key=lambda x: (x[0], x[1]))
  
  return predictions