409 lines
13 KiB
Python
409 lines
13 KiB
Python
import abc
|
|
import itertools
|
|
from collections import deque, defaultdict
|
|
import re
|
|
from typing import List, Optional, Dict, Any, Set, TypeVar
|
|
from dataclasses import dataclass
|
|
import networkx as nx
|
|
import penman
|
|
|
|
|
|
@dataclass
|
|
class SemanticGraph:
|
|
nodes_var: List[str]
|
|
"""
|
|
List of linearized nodes, with special tokens.
|
|
"""
|
|
edges: Optional[List[str]]
|
|
"""
|
|
List of linearized edges, with special tokens.
|
|
"""
|
|
backreferences: List[int]
|
|
"""
|
|
List of backpointers to handle rentrancies and cycles.
|
|
"""
|
|
var2instance: Dict[str, str]
|
|
"""
|
|
Dict from var ids to 'lemmatized' readable strings qualifying the node (collapsing the :instance edge for AMR).
|
|
"""
|
|
extra: Dict[str, Any]
|
|
"""
|
|
Holds extra stuff that might be useful, e.g. alignments, NER, EL.
|
|
"""
|
|
|
|
# @cached_property
|
|
@property
|
|
def variables(self) -> Set[str]:
|
|
"""Set of variables in this semantic graph"""
|
|
variables = {v for v in self.nodes_var if not v.startswith('<')}
|
|
return variables
|
|
|
|
@property
|
|
def resolved_nodes_var(self) -> List[str]:
|
|
return [self.nodes_var[b] for b in self.backreferences]
|
|
|
|
# @cached_property
|
|
@property
|
|
def nodes(self) -> List[str]:
|
|
"""Linearized nodes with varids replaced by instances"""
|
|
return [self.var2instance.get(node, node) for node in self.nodes_var]
|
|
|
|
@property
|
|
def resolved_nodes(self) -> List[str]:
|
|
return [self.nodes[b] for b in self.backreferences]
|
|
|
|
def src_occurrence(self, var: str) -> int:
|
|
pass
|
|
|
|
|
|
class BaseLinearizer(metaclass=abc.ABCMeta):
|
|
|
|
@abc.abstractmethod
|
|
def linearize(self, *args, **kwargs) -> SemanticGraph:
|
|
pass
|
|
|
|
|
|
class AMRTokens:
|
|
START, END = '<', '>'
|
|
_TEMPL = START + '{}' + END
|
|
|
|
BOS_N = _TEMPL.format('s')
|
|
EOS_N = _TEMPL.format('/s')
|
|
START_N = _TEMPL.format('start')
|
|
STOP_N = _TEMPL.format('stop')
|
|
PNTR_N = _TEMPL.format('pointer')
|
|
|
|
LIT_START = _TEMPL.format('lit')
|
|
LIT_END = _TEMPL.format('/lit')
|
|
|
|
BACKR_SRC_N = _TEMPL.format('backr:src:XXX')
|
|
BACKR_TRG_N = _TEMPL.format('backr:trg:XXX')
|
|
|
|
BOS_E = _TEMPL.format('s')
|
|
EOS_E = _TEMPL.format('/s')
|
|
START_E = _TEMPL.format('start')
|
|
STOP_E = _TEMPL.format('stop')
|
|
|
|
_FIXED_SPECIAL_TOKENS_N = {
|
|
BOS_N, EOS_N, START_N, STOP_N}
|
|
_FIXED_SPECIAL_TOKENS_E = {
|
|
BOS_E, EOS_E, START_E, STOP_E}
|
|
_FIXED_SPECIAL_TOKENS = _FIXED_SPECIAL_TOKENS_N | _FIXED_SPECIAL_TOKENS_E
|
|
|
|
# match and read backreferences
|
|
_re_BACKR_SRC_N = re.compile(BACKR_SRC_N.replace('XXX', r'([0-9]+)'))
|
|
_re_BACKR_TRG_N = re.compile(BACKR_TRG_N.replace('XXX', r'([0-9]+)'))
|
|
|
|
@classmethod
|
|
def is_node(cls, string: str) -> bool:
|
|
if isinstance(string, str) and string.startswith(':'):
|
|
return False
|
|
elif string in cls._FIXED_SPECIAL_TOKENS_E:
|
|
return False
|
|
return True
|
|
|
|
@classmethod
|
|
def read_backr(cls, string: str) -> Optional:
|
|
m_src = cls._re_BACKR_SRC_N.search(string)
|
|
if m_src is not None:
|
|
return m_src
|
|
m_trg = cls._re_BACKR_TRG_N.search(string)
|
|
if m_trg is not None:
|
|
return m_trg
|
|
return None
|
|
|
|
|
|
T = TypeVar('T')
|
|
|
|
|
|
def index_default(
|
|
item: T, list_: List[T],
|
|
start: Optional[int] = None,
|
|
stop: Optional[int] = None,
|
|
default: Optional[int] = None
|
|
):
|
|
if start is None:
|
|
start = 0
|
|
if stop is None:
|
|
stop = len(list_)
|
|
return next((i for i, x in enumerate(list_[start:stop], start=start) if x == item), default)
|
|
|
|
|
|
class AMRLinearizer(BaseLinearizer):
|
|
|
|
def __init__(
|
|
self,
|
|
use_pointer_tokens: bool = True,
|
|
collapse_name_ops: bool = False,
|
|
):
|
|
self.collapse_name_ops = collapse_name_ops
|
|
self.interleave_edges = False
|
|
self.use_pointer_tokens = use_pointer_tokens
|
|
|
|
def _collapse_name_ops(self, amr):
|
|
# identify name triples
|
|
name_vars = {}
|
|
for i, (v1, rel, v2) in enumerate(amr.triples):
|
|
if rel == ':instance' and v2 == 'name':
|
|
name_vars[v1] = 1
|
|
|
|
# check if they have ops
|
|
name_vars_to_ops = defaultdict(list)
|
|
for i, (v1, rel, v2) in enumerate(amr.triples):
|
|
if v1 in name_vars and rel.startswith(':op'):
|
|
name_vars_to_ops[v1].append((i, rel, v2.strip('"')))
|
|
|
|
triples = amr.triples.copy()
|
|
for nv, ops in name_vars_to_ops.items():
|
|
ops = sorted(ops, key=lambda x: int(x[1][3:]))
|
|
idx, _, lits = zip(*ops)
|
|
for i in idx:
|
|
triples[i] = None
|
|
lit = '"' + '_'.join(lits) + '"'
|
|
triples[min(idx)] = penman.Triple(nv, ':op1', lit)
|
|
|
|
triples = [t for t in triples if t is not None]
|
|
amr_ = penman.Graph(triples)
|
|
amr_.metadata = amr.metadata
|
|
return amr_
|
|
|
|
def linearize(self, amr: penman.Graph) -> SemanticGraph:
|
|
if self.collapse_name_ops:
|
|
amr = self._collapse_name_ops(amr)
|
|
linearized = self._linearize(amr)
|
|
linearized = self._interleave(linearized)
|
|
if self.use_pointer_tokens:
|
|
linearized = self._add_pointer_tokens(linearized)
|
|
return linearized
|
|
|
|
def _linearize(self, amr: penman.Graph) -> SemanticGraph:
|
|
variables = set(amr.variables())
|
|
variables = {'var:' + v for v in variables}
|
|
var2instance = {}
|
|
|
|
graph = nx.MultiDiGraph()
|
|
|
|
triples2order = {k: i for i, k in enumerate(amr.triples)}
|
|
|
|
for triple in amr.triples:
|
|
var, rel, instance = triple
|
|
order = triples2order[triple]
|
|
if rel != ':instance':
|
|
continue
|
|
for expansion_candidate in itertools.chain(range(order - 1, -1), range(order + 1, len(amr.triples))):
|
|
if var == amr.triples[expansion_candidate][2]:
|
|
expansion = expansion_candidate
|
|
break
|
|
else:
|
|
expansion = 0
|
|
var = 'var:' + var
|
|
var2instance[var] = instance
|
|
graph.add_node(var, instance=instance, order=order, expansion=expansion)
|
|
|
|
for triple in amr.edges():
|
|
var1, rel, var2 = triple
|
|
order = triples2order[triple]
|
|
if rel == ':instance':
|
|
continue
|
|
var1 = 'var:' + var1
|
|
var2 = 'var:' + var2
|
|
graph.add_edge(var1, var2, rel=rel, order=order)
|
|
|
|
for triple in amr.attributes():
|
|
var, rel, attr = triple
|
|
order = triples2order[triple]
|
|
if rel == ':instance':
|
|
continue
|
|
var = 'var:' + var
|
|
graph.add_edge(var, attr, rel=rel, order=order)
|
|
|
|
# nodes that are not reachable from the root (e.g. because of reification)
|
|
# will be present in the not_explored queue
|
|
# undirected_graph = graph.to_undirected()
|
|
# print(amr.variables())
|
|
not_explored = deque(sorted(variables, key=lambda x: nx.get_node_attributes(graph, 'order')[x]))
|
|
# (
|
|
# len(nx.shortest_path(undirected_graph, 'var:' + amr.top, x)),
|
|
# -graph.out_degree(x),
|
|
# )
|
|
|
|
first_index = {}
|
|
explored = set()
|
|
added_to_queue = set()
|
|
nodes_visit = [AMRTokens.BOS_N]
|
|
edges_visit = [AMRTokens.BOS_E]
|
|
backreferences = [0]
|
|
queue = deque()
|
|
queue.append('var:' + amr.top)
|
|
|
|
while queue or not_explored:
|
|
|
|
if queue:
|
|
node1 = queue.popleft()
|
|
else:
|
|
node1 = not_explored.popleft()
|
|
if node1 in added_to_queue:
|
|
continue
|
|
if not list(graph.successors(node1)):
|
|
continue
|
|
|
|
if node1 in variables:
|
|
if node1 in explored:
|
|
continue
|
|
if node1 in first_index:
|
|
nodes_visit.append(AMRTokens.BACKR_TRG_N)
|
|
backreferences.append(first_index[node1])
|
|
else:
|
|
backreferences.append(len(nodes_visit))
|
|
first_index[node1] = len(nodes_visit)
|
|
nodes_visit.append(node1)
|
|
edges_visit.append(AMRTokens.START_E)
|
|
|
|
successors = []
|
|
for node2 in graph.successors(node1):
|
|
for edge_data in graph.get_edge_data(node1, node2).values():
|
|
rel = edge_data['rel']
|
|
order = edge_data['order']
|
|
successors.append((order, rel, node2))
|
|
successors = sorted(successors)
|
|
|
|
for order, rel, node2 in successors:
|
|
edges_visit.append(rel)
|
|
|
|
# node2 is a variable
|
|
if node2 in variables:
|
|
# ... which was mentioned before
|
|
if node2 in first_index:
|
|
nodes_visit.append(AMRTokens.BACKR_TRG_N)
|
|
backreferences.append(first_index[node2])
|
|
|
|
# .. which is mentioned for the first time
|
|
else:
|
|
backreferences.append(len(nodes_visit))
|
|
first_index[node2] = len(nodes_visit)
|
|
nodes_visit.append(node2)
|
|
|
|
# 1) not already in Q
|
|
# 2) has children
|
|
# 3) the edge right before its expansion has been encountered
|
|
if (node2 not in added_to_queue) and list(graph.successors(node2)) and (
|
|
nx.get_node_attributes(graph, 'expansion')[node2] <= order):
|
|
queue.append(node2)
|
|
added_to_queue.add(node2)
|
|
|
|
# node2 is a constant
|
|
else:
|
|
backreferences.append(len(nodes_visit))
|
|
nodes_visit.append(node2)
|
|
|
|
backreferences.append(len(nodes_visit))
|
|
nodes_visit.append(AMRTokens.STOP_N)
|
|
edges_visit.append(AMRTokens.STOP_E)
|
|
explored.add(node1)
|
|
|
|
else:
|
|
backreferences.append(len(nodes_visit))
|
|
nodes_visit.append(node1)
|
|
explored.add(node1)
|
|
|
|
backreferences.append(len(nodes_visit))
|
|
nodes_visit.append(AMRTokens.EOS_N)
|
|
edges_visit.append(AMRTokens.EOS_E)
|
|
assert len(nodes_visit) == len(edges_visit) == len(backreferences)
|
|
return SemanticGraph(
|
|
nodes_visit,
|
|
edges_visit,
|
|
backreferences,
|
|
var2instance,
|
|
extra={'graph': graph, 'amr': amr}
|
|
)
|
|
|
|
def _interleave(self, graph: SemanticGraph) -> SemanticGraph:
|
|
|
|
new_backreferences_map = []
|
|
new_nodes = []
|
|
new_edges = None
|
|
new_backreferences = []
|
|
|
|
# to isolate sublist to the stop token
|
|
start_i = 1
|
|
end_i = index_default(AMRTokens.STOP_N, graph.nodes_var, start_i, -1, -1)
|
|
|
|
def add_node(node, backr=None):
|
|
old_n_node = len(new_backreferences_map)
|
|
new_n_node = len(new_nodes)
|
|
|
|
if backr is None:
|
|
backr = old_n_node
|
|
|
|
new_backreferences_map.append(new_n_node)
|
|
new_nodes.append(node)
|
|
if old_n_node == backr:
|
|
new_backreferences.append(new_n_node)
|
|
else:
|
|
new_backreferences.append(new_backreferences_map[backr])
|
|
|
|
def add_edge(edge):
|
|
new_nodes.append(edge)
|
|
new_backreferences.append(len(new_backreferences))
|
|
|
|
add_node(AMRTokens.BOS_N)
|
|
|
|
while end_i > -1:
|
|
|
|
# src node
|
|
add_node(graph.nodes_var[start_i], graph.backreferences[start_i])
|
|
|
|
# edges and trg nodes, interleaved
|
|
nodes = graph.nodes_var[start_i + 1:end_i]
|
|
edges = graph.edges[start_i + 1:end_i]
|
|
backr = graph.backreferences[start_i + 1:end_i]
|
|
for n, e, b in zip(nodes, edges, backr):
|
|
add_edge(e)
|
|
add_node(n, b)
|
|
|
|
# stop
|
|
add_node(graph.nodes_var[end_i], graph.backreferences[end_i])
|
|
|
|
start_i = end_i + 1
|
|
end_i = index_default(AMRTokens.STOP_N, graph.nodes_var, start_i, -1, -1)
|
|
|
|
add_node(AMRTokens.EOS_N)
|
|
|
|
new_graph = SemanticGraph(
|
|
new_nodes,
|
|
None,
|
|
new_backreferences,
|
|
graph.var2instance,
|
|
extra=graph.extra,
|
|
)
|
|
return new_graph
|
|
|
|
def _add_pointer_tokens(self, graph: SemanticGraph) -> SemanticGraph:
|
|
new_nodes = []
|
|
var2pointer = {}
|
|
for node, backr in zip(graph.nodes_var, graph.backreferences):
|
|
|
|
if node == AMRTokens.BACKR_TRG_N:
|
|
node = graph.nodes_var[backr]
|
|
pointer = var2pointer[node]
|
|
new_nodes.append(pointer)
|
|
elif node in graph.var2instance:
|
|
pointer = var2pointer.setdefault(node, f"<pointer:{len(var2pointer)}>")
|
|
new_nodes.append(pointer)
|
|
new_nodes.append(node)
|
|
else:
|
|
new_nodes.append(node)
|
|
|
|
new_backreferences = list(range(len(new_nodes)))
|
|
new_graph = SemanticGraph(
|
|
new_nodes,
|
|
None,
|
|
new_backreferences,
|
|
graph.var2instance,
|
|
extra=graph.extra,
|
|
)
|
|
return new_graph
|
|
|
|
|