195 lines
6.3 KiB
Python
195 lines
6.3 KiB
Python
"""Encode English text, commit supported graph edits, and decode a revision."""
|
|
|
|
import argparse
|
|
import json
|
|
from dataclasses import dataclass, replace
|
|
from graphlib import TopologicalSorter
|
|
|
|
import spacy
|
|
import tiktoken
|
|
|
|
|
|
MODEL_ENCODING = "o200k_base"
|
|
|
|
|
|
@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, ...] = ()
|
|
input_ids: tuple[int, ...] = ()
|
|
|
|
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.tokenizer = tiktoken.get_encoding(MODEL_ENCODING)
|
|
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:
|
|
content = [token.i for token in sentence if not token.is_space]
|
|
if not content:
|
|
continue
|
|
sentence = doc[content[0]:content[-1] + 1]
|
|
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),
|
|
input_ids=tuple(self.tokenizer.encode_ordinary(text)),
|
|
)
|
|
|
|
@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_ not in {"VBD", "VBN"}
|
|
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 if not token.is_punct}
|
|
!= set(range(marker.i, end))
|
|
):
|
|
return None
|
|
return marker, end
|
|
|
|
|
|
class Decoder:
|
|
def decode_ids(self, graph: SyntaxDAG) -> tuple[int, ...]:
|
|
tokenizer = tiktoken.get_encoding(MODEL_ENCODING)
|
|
return tuple(tokenizer.encode_ordinary(self.decode(graph)))
|
|
|
|
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")
|
|
parser.add_argument("--ids", action="store_true", help="Emit model token IDs as JSON")
|
|
args = parser.parse_args()
|
|
graph = Encoder().encode(args.text)
|
|
if args.commit:
|
|
for proposal in graph.proposals:
|
|
graph = graph.commit(proposal)
|
|
decoder = Decoder()
|
|
text = decoder.decode(graph)
|
|
if args.ids:
|
|
print(json.dumps({
|
|
"encoding": MODEL_ENCODING,
|
|
"input_ids": graph.input_ids,
|
|
"output_ids": decoder.decode_ids(graph),
|
|
"text": text,
|
|
}, ensure_ascii=False))
|
|
else:
|
|
print(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|