File size: 2,132 Bytes
1d5604f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
#!/usr/bin/env python3
# coding=utf-8
from data.parser.to_mrp.abstract_parser import AbstractParser
class LabeledEdgeParser(AbstractParser):
def __init__(self, *args):
super().__init__(*args)
self.source_id = self.dataset.edge_label_field.vocab.stoi["Source"]
self.target_id = self.dataset.edge_label_field.vocab.stoi["Target"]
def parse(self, prediction):
output = {}
output["id"] = self.dataset.id_field.vocab.itos[prediction["id"].item()]
output["nodes"] = self.create_nodes(prediction)
output["nodes"] = self.create_anchors(prediction, output["nodes"], join_contiguous=True, at_least_one=True)
output["nodes"] = [{"id": 0}] + output["nodes"]
output["edges"] = self.create_edges(prediction, output["nodes"])
return output
def create_nodes(self, prediction):
return [{"id": i + 1} for i, l in enumerate(prediction["labels"])]
def create_edges(self, prediction, nodes):
N = len(nodes)
edge_prediction = prediction["edge presence"][:N, :N]
edges = []
for target in range(1, N):
if edge_prediction[0, target] >= 0.5:
prediction["edge labels"][0, target, self.source_id] = float("-inf")
prediction["edge labels"][0, target, self.target_id] = float("-inf")
self.create_edge(0, target, prediction, edges, nodes)
for source in range(1, N):
for target in range(1, N):
if source == target:
continue
if edge_prediction[source, target] < 0.5:
continue
for i in range(prediction["edge labels"].size(2)):
if i not in [self.source_id, self.target_id]:
prediction["edge labels"][source, target, i] = float("-inf")
self.create_edge(source, target, prediction, edges, nodes)
return edges
def get_edge_label(self, prediction, source, target):
return self.dataset.edge_label_field.vocab.itos[prediction["edge labels"][source, target].argmax(-1).item()]
|