baim/synthetic.py

"""Seeded synthetic supervision. Templates are fixtures, never runtime rules."""
import argparse
from dataclasses import asdict
import html
import json
from pathlib import Path
import random
from .state import Element, digest

TRAIN_TEMPLATES = ('stack', 'grid', 'fieldset')
TEST_TEMPLATES = ('table', 'nested')
WORDS = ('amber birch cobalt delta elm fern granite harbor indigo jade kelp linen '
         'maple nectar olive pearl quartz river silver timber umber violet willow xenon yellow zinc').split()
TRAIN_WORDING = {
    'C': ['Click {name}.', 'Activate {name}.', 'Choose the control named {name}.'],
    'T': ['Enter "{value}" into {name}.', 'Fill {name} with "{value}".', 'Type "{value}" in {name}.'],
    'O': ['Select "{value}" from {name}.', 'Set the dropdown {name} to "{value}".', 'Choose "{value}" in the {name} list.'],
}
NOVEL_WORDING = {
    'C': ['Press the {name} control.', 'Use the item labelled {name}.'],
    'T': ['Populate {name} with "{value}".', 'Put "{value}" into the field labelled {name}.'],
    'O': ['Pick "{value}" from the menu labelled {name}.', 'Change the {name} selector to "{value}".'],
}


def generate(split, count, seed, novel_wording=False):
    rng = random.Random(seed)
    templates = TEST_TEMPLATES if split == 'test' else TRAIN_TEMPLATES
    for sample_id in range(count):
        kind = rng.choice(('C', 'T', 'O'))
        total = rng.randint(8, 32)
        names = rng.sample([f'{a} {b}' for a in WORDS for b in WORDS if a != b], total)
        target = rng.randrange(total)
        rows = []
        for index, name in enumerate(names):
            role = rng.choice(('button', 'textbox', 'combobox', 'link'))
            if index == target:
                role = rng.choice(('button', 'link')) if kind == 'C' else {'T':'textbox','O':'combobox'}[kind]
            rows.append(asdict(Element(f'e{index}', role, name, str(index))))
        value = f'value-{rng.randrange(1000000)}'
        wording = NOVEL_WORDING if novel_wording else TRAIN_WORDING
        goal = rng.choice(wording[kind]).format(name=names[target], value=value)
        yield dict(schema_version=1, source='baim-synthetic-v1', license='CC0-1.0',
                   split=split, template=rng.choice(templates), seed=seed, sample_id=sample_id,
                   goal=goal, elements=rows, action=kind, target=target,
                   argument=value if kind in {'T','O'} else None,
                   execution_verified=False, content_hash=digest([goal, rows]))


def render(sample):
    """Returns HTML plus an independent DOM-side outcome oracle for test tasks."""
    controls = []
    for index, row in enumerate(sample['elements']):
        name = html.escape(row['name'], quote=True)
        ident = f'control-{sample["seed"]}-{sample["sample_id"]}-{index}'
        attrs = f'id="{ident}" data-index="{index}" aria-label="{name}"'
        role = row['role']
        if role == 'textbox':
            control = f'<input {attrs}>'
        elif role == 'combobox':
            value = html.escape(sample['argument'] or 'option', quote=True)
            control = f'<select {attrs}><option value="">Unset</option><option value="{value}">{value}</option></select>'
        elif role == 'link':
            control = f'<a {attrs} href="#result" onclick="window.fixtureResult={index}">{name}</a>'
        else:
            control = f'<button {attrs} onclick="window.fixtureResult={index}">{name}</button>'
        controls.append(control)
    template = sample['template']
    if template == 'table':
        body = '<table>' + ''.join(f'<tr><td>{x}</td></tr>' for x in controls) + '</table>'
    elif template == 'nested':
        body = '<article>' + ''.join(f'<section><div><aside>{x}</aside></div></section>' for x in controls) + '</article>'
    elif template == 'fieldset':
        body = '<form>' + ''.join(f'<fieldset>{x}</fieldset>' for x in controls) + '</form>'
    elif template == 'grid':
        body = '<main style="display:grid;grid-template-columns:repeat(3,1fr)">' + ''.join(controls) + '</main>'
    else:
        body = '<main>' + ''.join(f'<div>{x}</div>' for x in controls) + '</main>'
    return '<!doctype html><meta charset="utf-8"><style>input,button,select,a{margin:6px;padding:5px}</style>' + body


def load(path):
    return [json.loads(line) for line in Path(path).read_text(encoding='utf-8').splitlines() if line]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default='datasets/synthetic-v1')
    args = parser.parse_args()
    root = Path(args.output)
    root.mkdir(parents=True, exist_ok=True)
    manifest = {}
    for split, count, seed, novel in [('train',2400,101,False),('validation',480,202,False),
                                      ('test',480,303,False),('novel_wording',480,404,True)]:
        rows = list(generate('test' if split == 'novel_wording' else split, count, seed, novel))
        data = ''.join(json.dumps(row, ensure_ascii=False) + '\n' for row in rows)
        (root / f'{split}.jsonl').write_text(data, encoding='utf-8')
        manifest[split] = dict(count=count, seed=seed, templates=sorted({r['template'] for r in rows}),
                               sha256=digest(data))
    manifest['limitations'] = 'Single-step generated tasks; test layouts held out but vocabulary shared. No arbitrary-site claim.'
    (root / 'manifest.json').write_text(json.dumps(manifest, indent=2), encoding='utf-8')
    print(json.dumps(manifest, indent=2))


if __name__ == '__main__':
    main()