wehub-resource-sync
244 lines
10 KiB
Python
244 lines
10 KiB
Python
# Adopted from https://github.com/KiroSummer/A_Syntax-aware_MTL_Framework_for_Chinese_SRL
|
|
|
|
# Inference functions for the SRL model.
|
|
import numpy as np
|
|
|
|
|
|
def decode_spans(span_starts, span_ends, span_scores, labels_inv):
|
|
"""
|
|
|
|
Args:
|
|
span_starts: [num_candidates,]
|
|
span_scores: [num_candidates, num_labels]
|
|
span_ends:
|
|
labels_inv:
|
|
|
|
Returns:
|
|
|
|
|
|
"""
|
|
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:
|
|
|
|
|
|
"""
|
|
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 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] # [num_arg, num_roles]
|
|
labels = np.argmax(scores, axis=1).astype(np.int64)
|
|
spans = list(zip(starts, ends, list(range(len(starts)))))
|
|
spans = sorted(spans, key=lambda x: (x[0], x[1])) # sort according to the span start index
|
|
|
|
if u_constraint:
|
|
f = np.zeros([max_len + 1, 128], dtype=float) - 0.1
|
|
else: # This one
|
|
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) # the pointers store
|
|
if f[t1][rs1] > f[best_state[0][0]][best_state[0][1]]:
|
|
best_state[0] = (t1, rs1)
|
|
|
|
for start, end, i in spans: # [arg_start, arg_end, arg_span_id]
|
|
assert scores[i][0] == 0 # dummy score
|
|
# 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: # skip the span contains the predicate
|
|
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: # labels_inv[r0] == "O"
|
|
continue
|
|
r0_str = labels_inv[r0]
|
|
# Enumerate explored states.
|
|
t_states = [t for t in list(states.keys()) if t <= start] # collect the state which is before the current span
|
|
for t in t_states: # for each state
|
|
role_states = states[t]
|
|
# Update states if best role is not a core arg.
|
|
if not u_constraint or r0_str not in _CORE_ARGS: # True; this one
|
|
for rs in role_states: # the set type in the value in the state dict
|
|
_update_state(t, rs, end + 1, rs, scores[i][r0], i, r0) # update the state
|
|
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((int(starts[i]), int(ends[i]), labels_inv[r]))
|
|
t = t0
|
|
rs = rs0
|
|
return new_spans[::-1]
|
|
|
|
|
|
def srl_decode(sentence_lengths, predict_dict, srl_labels_inv, config): # decode the predictions.
|
|
# Decode sentence-level tasks.
|
|
num_sentences = len(sentence_lengths)
|
|
predictions = [{} for _ in range(num_sentences)]
|
|
# Sentence-level predictions.
|
|
for i in range(num_sentences): # for each sentences
|
|
# if predict_dict["No_arg"] is True:
|
|
# predictions["srl"][i][predict_dict["predicates"][i]] = []
|
|
# continue
|
|
predict_dict_num_args_ = predict_dict["num_args"].cpu().numpy()
|
|
predict_dict_num_preds_ = predict_dict["num_preds"].cpu().numpy()
|
|
predict_dict_predicates_ = predict_dict["predicates"].cpu().numpy()
|
|
predict_dict_arg_starts_ = predict_dict["arg_starts"].cpu().numpy()
|
|
predict_dict_arg_ends_ = predict_dict["arg_ends"].cpu().numpy()
|
|
predict_dict_srl_scores_ = predict_dict["srl_scores"].detach().cpu().numpy()
|
|
num_args = predict_dict_num_args_[i] # the number of the candidate argument spans
|
|
num_preds = predict_dict_num_preds_[i] # the number of the candidate predicates
|
|
# for each predicate id, exec the decode process
|
|
for j, pred_id in enumerate(predict_dict_predicates_[i][:num_preds]):
|
|
# sorted arg_starts and arg_ends and srl_scores ? should be??? enforce_srl_constraint = False
|
|
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, :],
|
|
sentence_lengths[i], srl_labels_inv, pred_id, config.enforce_srl_constraint)
|
|
# To avoid warnings in the eval script.
|
|
if config.use_gold_predicates: # false
|
|
arg_spans.append((pred_id, pred_id, "V"))
|
|
if arg_spans:
|
|
predictions[i][int(pred_id)] = sorted(arg_spans, key=lambda x: (x[0], x[1]))
|
|
|
|
return predictions
|