79 lines
3.2 KiB
Python
79 lines
3.2 KiB
Python
import unittest
|
|
from dataclasses import replace
|
|
from graphlib import CycleError
|
|
|
|
import tiktoken
|
|
|
|
from syntax_dag import Decoder, Edge, Encoder, Proposal
|
|
|
|
|
|
class SyntaxDAGTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.encoder = Encoder()
|
|
cls.decoder = Decoder()
|
|
|
|
def test_round_trip(self):
|
|
for text in ("", " \t", " It rained.\n\nThe road flooded. "):
|
|
with self.subTest(text=text):
|
|
graph = self.encoder.encode(text)
|
|
self.assertEqual(self.decoder.decode(graph), text)
|
|
self.assertFalse(graph.proposals)
|
|
|
|
def test_causal_commit(self):
|
|
text = "The road flooded because it rained. They waited."
|
|
graph = self.encoder.encode(text)
|
|
self.assertEqual(len(graph.proposals), 1, [
|
|
(token.text, token.head.text, token.dep_, token.pos_, token.tag_)
|
|
for token in self.encoder.nlp(text)
|
|
])
|
|
proposal, = graph.proposals
|
|
self.assertTrue(any(
|
|
len({edge.source for edge in graph.edges if edge.target == node}) > 1
|
|
for node in range(len(graph.tokens))
|
|
))
|
|
committed = graph.commit(proposal)
|
|
self.assertEqual(self.decoder.decode(graph), text)
|
|
self.assertEqual(
|
|
self.decoder.decode(committed),
|
|
"Cause: it rained. Effect: The road flooded. They waited.",
|
|
)
|
|
self.assertEqual(committed.history, (proposal,))
|
|
self.assertIs(committed.commit(proposal), committed)
|
|
self.assertNotIn(proposal.edge, graph.edges)
|
|
|
|
def test_cycle_rejects_entire_commit(self):
|
|
graph = self.encoder.encode("The road flooded because it rained.")
|
|
original, = graph.proposals
|
|
reverse = Proposal(Edge(original.edge.target, original.edge.source, "causes"),
|
|
original.marker, "test")
|
|
committed = replace(graph.commit(original), proposals=(original, reverse))
|
|
with self.assertRaises(CycleError):
|
|
committed.commit(reverse)
|
|
self.assertEqual(committed.history, (original,))
|
|
self.assertNotIn(reverse.edge, committed.edges)
|
|
|
|
def test_negation_and_questions_are_left_alone(self):
|
|
for text in ("The road did not flood because it rained.",
|
|
"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),
|
|
)
|