Prototype causal rewriting with a syntax DAG

syntax_dag.py

@@ -0,0 +1,164 @@
+"""Encode English text, commit supported graph edits, and decode a revision."""
+
+import argparse
+from dataclasses import dataclass, replace
+from graphlib import TopologicalSorter
+
+import spacy
+
+
+@dataclass(frozen=True)
+class Token:
+    text: str
+    whitespace: str
+
+
+@dataclass(frozen=True)
+class Clause:
+    id: int
+    start: int
+    stop: int
+
+
+@dataclass(frozen=True)
+class Edge:
+    source: int
+    target: int
+    kind: str
+
+
+@dataclass(frozen=True)
+class Proposal:
+    edge: Edge
+    marker: int
+    rule: str
+
+
+@dataclass(frozen=True)
+class SyntaxDAG:
+    tokens: tuple[Token, ...]
+    clauses: tuple[Clause, ...]
+    edges: tuple[Edge, ...]
+    proposals: tuple[Proposal, ...]
+    history: tuple[Proposal, ...] = ()
+
+    def __post_init__(self):
+        nodes = set(range(len(self.tokens))) | {clause.id for clause in self.clauses}
+        sorter = TopologicalSorter()
+        for node in nodes:
+            sorter.add(node)
+        for edge in self.edges:
+            if edge.source not in nodes or edge.target not in nodes:
+                raise ValueError("Edge endpoint is missing")
+            sorter.add(edge.target, edge.source)
+        tuple(sorter.static_order())
+
+    def commit(self, proposal: Proposal):
+        if proposal not in self.proposals:
+            raise ValueError("Proposal was not mined from this graph")
+        if proposal in self.history:
+            return self
+        return replace(
+            self,
+            edges=(*self.edges, proposal.edge),
+            history=(*self.history, proposal),
+        )
+
+
+class Encoder:
+    def __init__(self):
+        self.nlp = spacy.load("en_core_web_sm", disable=["ner"])
+
+    def encode(self, text: str) -> SyntaxDAG:
+        doc = self.nlp(text)
+        tokens = tuple(Token(token.text, token.whitespace_) for token in doc)
+        edges = [
+            Edge(token.head.i, token.i, token.dep_)
+            for token in doc
+            if token.head != token
+        ]
+        clauses = []
+        proposals = []
+        for sentence in doc.sents:
+            match = self._causal_marker(sentence)
+            if match is None:
+                continue
+            marker, end = match
+            effect = Clause(len(tokens) + len(clauses), sentence.start, marker.i)
+            cause = Clause(effect.id + 1, marker.i + 1, end)
+            clauses.extend((effect, cause))
+            for clause in (effect, cause):
+                edges.extend(
+                    Edge(clause.id, index, "contains")
+                    for index in range(clause.start, clause.stop)
+                )
+            proposals.append(Proposal(Edge(cause.id, effect.id, "causes"),
+                                      marker.i, "trailing-because"))
+        return SyntaxDAG(tokens, tuple(clauses), tuple(edges), tuple(proposals))
+
+    @staticmethod
+    def _causal_marker(sentence):
+        markers = [token for token in sentence if token.lower_ == "because"]
+        if len(markers) != 1 or sentence[-1].text != ".":
+            return None
+        marker = markers[0]
+        verb = marker.head
+        end = sentence.end - 1
+        if (
+            marker.dep_ != "mark"
+            or verb.dep_ != "advcl"
+            or verb.head != sentence.root
+            or not sentence.start <= sentence.root.i < marker.i
+            or not marker.i < verb.i < end
+            or sentence.root.pos_ != "VERB"
+            or verb.pos_ != "VERB"
+            or sentence.root.tag_ != "VBD"
+            or verb.tag_ != "VBD"
+            or any(token.is_punct or token.is_space for token in sentence[:-1])
+            or any(token.dep_ in {"neg", "aux", "conj"} for token in sentence)
+            or not any(child.dep_ == "nsubj" for child in sentence.root.children)
+            or not any(child.dep_ == "nsubj" for child in verb.children)
+            or {token.i for token in verb.subtree} != set(range(marker.i, end))
+        ):
+            return None
+        return marker, end
+
+
+class Decoder:
+    def decode(self, graph: SyntaxDAG) -> str:
+        clauses = {clause.id: clause for clause in graph.clauses}
+        edits = {}
+        for edge in graph.edges:
+            if edge.kind != "causes":
+                continue
+            cause, effect = clauses[edge.source], clauses[edge.target]
+            cause_text = self._surface(graph, cause.start, cause.stop).rstrip()
+            effect_text = self._surface(graph, effect.start, effect.stop).rstrip()
+            edits[effect.start] = (cause.stop, f"Cause: {cause_text}. Effect: {effect_text}")
+        output = []
+        index = 0
+        while index < len(graph.tokens):
+            stop, surface = edits.get(index, (index + 1, self._surface(graph, index, index + 1)))
+            output.append(surface)
+            index = stop
+        return "".join(output)
+
+    @staticmethod
+    def _surface(graph, start, stop):
+        return "".join(token.text + token.whitespace for token in graph.tokens[start:stop])
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("text")
+    parser.add_argument("--commit", action="store_true", help="Apply mined causal edits")
+    args = parser.parse_args()
+    graph = Encoder().encode(args.text)
+    if args.commit:
+        for proposal in graph.proposals:
+            graph = graph.commit(proposal)
+    print(Decoder().decode(graph))
+
+
+if __name__ == "__main__":
+    main()