Merge pull request 'Integrate the syntax DAG and model token IDs' (#3) from codex/dag-tokenizer-integration into main

Reviewed-on: #3

README.md

@@ -34,3 +34,16 @@ The first rule covers simple affirmative past-tense statements with a trailing
 finite causal clause. Extraction is heuristic and depends on parser accuracy.
 Commits are in-memory graph edits. The prototype uses explicit authorization
 via the CLI flag and has no learned rewrite policy.
+
+## Model token IDs
+
+```sh
+python -m syntax_dag --ids --commit 'The road flooded because it rained.'
+```
+
+JSON includes `encoding`, `input_ids`, `output_ids`, and the output `text`.
+These are integer vocabulary IDs from tiktoken's `o200k_base` BPE encoding.
+spaCy tokens still serve as syntax nodes. Their positions are independent of
+BPE positions. A commit preserves the original input IDs, and the decoder
+tokenizes the full rendered text so its output IDs reflect the new boundaries.
+The tokenizer downloads its vocabulary on first use if it is absent from cache.

requirements.txt

@@ -3,3 +3,5 @@ en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_
 
 # spaCy 3.8.11 imports Click when loading its CLI.
 click==8.1.8
+
+tiktoken==0.12.0

syntax_dag.py

@@ -1,10 +1,15 @@
 """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)
@@ -41,6 +46,7 @@ class SyntaxDAG:
     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}
@@ -67,6 +73,7 @@ class SyntaxDAG:
 
 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:
@@ -98,7 +105,10 @@ class Encoder:
                 )
             proposals.append(Proposal(Edge(cause.id, effect.id, "causes"),
                                       marker.i, "trailing-because"))
-        return SyntaxDAG(tokens, tuple(clauses), tuple(edges), tuple(proposals))
+        return SyntaxDAG(
+            tokens, tuple(clauses), tuple(edges), tuple(proposals),
+            input_ids=tuple(self.tokenizer.encode_ordinary(text)),
+        )
 
     @staticmethod
     def _causal_marker(sentence):
@@ -130,6 +140,10 @@ class Encoder:
 
 
 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 = {}
@@ -157,12 +171,23 @@ 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)
-    print(Decoder().decode(graph))
+    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__":

test_syntax_dag.py

@@ -2,6 +2,8 @@ import unittest
 from dataclasses import replace
 from graphlib import CycleError
 
+import tiktoken
+
 from syntax_dag import Decoder, Edge, Encoder, Proposal
 
 
@@ -56,3 +58,21 @@ class SyntaxDAGTest(unittest.TestCase):
                      "Did the road flood because it rained?"):
             with self.subTest(text=text):
                 self.assertFalse(self.encoder.encode(text).proposals)
+
+
+    def test_model_token_ids(self):
+        tokenizer = tiktoken.get_encoding("o200k_base")
+        text = "中文 🙂 <|endoftext|>"
+        graph = self.encoder.encode(text)
+        self.assertEqual(graph.input_ids, tuple(tokenizer.encode_ordinary(text)))
+        self.assertEqual(tokenizer.decode(graph.input_ids), text)
+        self.assertEqual(self.decoder.decode_ids(graph), graph.input_ids)
+
+        graph = self.encoder.encode("The road flooded because it rained.")
+        committed = graph.commit(graph.proposals[0])
+        self.assertEqual(committed.input_ids, graph.input_ids)
+        self.assertNotEqual(self.decoder.decode_ids(committed), graph.input_ids)
+        self.assertEqual(
+            tokenizer.decode(self.decoder.decode_ids(committed)),
+            self.decoder.decode(committed),
+        )