Files
2026-07-13 12:37:18 +08:00

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