File size: 2,435 Bytes
1d26c54 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import sys, tempfile, unittest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from document_run import transform
import geometry
class FakeEngine:
identity = {'test': True}
progress = staticmethod(lambda message: None)
def __init__(self): self.calls = []
def generate(self, text, **options):
self.calls.append((text, options))
if len(self.calls) == 1: raise RuntimeError('Test generation failure')
return {'text': 'A dark uncertainty.'}
class DocumentRun(unittest.TestCase):
def test_failure_continues_and_settings_are_preserved(self):
engine = FakeEngine()
g = geometry.new()
g['controls'].update(strength=.47,budget=.47)
text = 'The machine predicts the next observation.\n\nIt revises its estimates when the inputs change.'
with tempfile.TemporaryDirectory() as tmp:
path = Path(tmp)/'run.json'
report = transform(engine,text,g,'Creative',dict(max_tokens=32,temperature=0.,top_p=.95,top_k=40,time_limit=120),path)
self.assertTrue(report['complete'])
self.assertEqual(report['failed'],1)
self.assertIn('Test generation failure',report['records']['0']['attempts'][0]['error'])
self.assertTrue(path.with_suffix('.txt').exists())
self.assertEqual(len(engine.calls),2)
for passage,opts in engine.calls:
self.assertEqual(opts['g']['controls'],g['controls'])
self.assertEqual(opts['max_chars'],int(len(passage)*1.25))
self.assertEqual(opts['soft_chars'],int(len(passage)*.85))
self.assertGreaterEqual(opts['max_tokens'],opts['max_chars'])
self.assertIn('Meaning, imagery, and events may change',opts['system'])
self.assertLess(report['records']['0']['target'][0],report['records']['1']['target'][0])
def test_chopped_ending_is_not_joined_to_next_paragraph(self):
engine=FakeEngine()
engine.generate=lambda text,**opts: {'text':'Dread gathers. Unfinish'}
with tempfile.TemporaryDirectory() as tmp:
report=transform(engine,'An ordinary and very quiet day.\n\nAnother ordinary and quiet day.',geometry.new(),'Restrained',dict(max_tokens=512),Path(tmp)/'run.json')
self.assertEqual(report['text'],'Dread gathers.\n\nDread gathers.')
self.assertIn('Unfinish',report['records']['0']['attempts'][0]['text'])
|