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 = 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 = list(zip(arg_starts, arg_ends, list(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 list(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 list(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 = list(zip(starts, ends, list(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 list(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

    # print spans
    # print new_spans[::-1]
    return new_spans[::-1]


# One-stop decoder for all the tasks.
def mtl_decode(sentences, predict_dict, ner_labels_inv, rel_labels_inv, config):
    predictions = {}

    # Decode sentence-level tasks.
    num_sentences = len(sentences)
    if "srl_scores" in predict_dict:
        predictions["srl"] = [{} for i in range(num_sentences)]
    if "ner_scores" in predict_dict:
        predictions["ner"] = [{} for i in range(num_sentences)]
    if "rel_scores" in predict_dict:
        predictions["rel"] = [[] for i in range(num_sentences)]

    # Sentence-level predictions.
    for i in range(num_sentences):
        if "rel" in predictions:

            # predictions['rel'][i].append(predict_dict['rel_scores'][i])
            # predictions['rel'][i].append(predict_dict['entity_starts'][i])
            # predictions['rel'][i].append(predict_dict['entity_ends'][i])
            # predictions['rel'][i].append(predict_dict["num_entities"][i])

            num_ents = predict_dict["num_entities"][i]
            ent_starts = predict_dict["entity_starts"][i]
            ent_ends = predict_dict["entity_ends"][i]
            for j in range(num_ents):
                for k in range(num_ents):
                    pred = predict_dict["rel_labels"][i, j, k]
                    if pred > 0:
                        predictions["rel"][i].append([
                            ent_starts[j], ent_ends[j], ent_starts[k], ent_ends[k],
                            rel_labels_inv[pred]])
        if "ner" in predictions:
            ner_spans = _dp_decode_non_overlapping_spans(
                predict_dict["candidate_starts"][i],
                predict_dict["candidate_ends"][i],
                predict_dict["ner_scores"][i],
                len(sentences[i]), ner_labels_inv, None, False)
            predictions["ner"][i] = ner_spans

    # Document-level predictions. -1 means null antecedent.
    if "antecedent_scores" in predict_dict:
        mention_spans = list(zip(predict_dict["mention_starts"], predict_dict["mention_ends"]))
        mention_to_predicted = {}
        predicted_clusters = []

        def _link_mentions(curr_span, ant_span):
            if ant_span not in mention_to_predicted:
                new_cluster_id = len(predicted_clusters)
                mention_to_predicted[ant_span] = new_cluster_id
                predicted_clusters.append([ant_span, ])
            cluster_id = mention_to_predicted[ant_span]
            if not curr_span in mention_to_predicted:
                mention_to_predicted[curr_span] = cluster_id
                predicted_clusters[cluster_id].append(curr_span)
            '''else:
              cluster_id2 = mention_to_predicted[curr_span]
              # Merge clusters.
              if cluster_id != cluster_id2:
                print "Merging clusters:", predicted_clusters[cluster_id], predicted_clusters[cluster_id2]
                for span in predicted_clusters[cluster_id2]:
                  mention_to_predicted[span] = cluster_id
                  predicted_clusters[cluster_id].append(span)
                predicted_clusters[cluster_id2] = []'''

        scores = predict_dict["antecedent_scores"]
        antecedents = predict_dict["antecedents"]
        # if config["coref_loss"] == "mention_rank":
        for i, ant_label in enumerate(np.argmax(scores, axis=1)):
            if ant_label <= 0:
                continue
            ant_id = antecedents[i][ant_label - 1]
            assert i > ant_id
            _link_mentions(mention_spans[i], mention_spans[ant_id])
        '''else:
          for i, curr_span in enumerate(mention_spans):
            for j in range(1, scores.shape[1]):
              if scores[i][j] > 0:
                _link_mentions(curr_span, mention_spans[antecedents[i][j-1]])'''

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

    # print predictions["srl"]
    return predictions