| """Public, generated supervision for the parameter-count experiment. |
| |
| The assembler supplies the exact changed byte (or NONE). This is supervised |
| training data; it is not an execution-feedback or native self-learning loop. |
| Independent final evaluation lives in a separate, unopened data directory. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| from collections import Counter |
| import hashlib |
| import json |
| from pathlib import Path |
| import random |
| import sys |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from lightnest import rv32 |
| from lightnest.learned_patch import TRAIN_REGS, TRAIN_TEMPLATES |
|
|
|
|
| NEGATION_TEMPLATES = ( |
| "do not change {a} to {b}", |
| "do not turn {a} into {b}", |
| "do not replace {a} with {b}", |
| "do not switch from {a} to {b}", |
| ) |
| FAMILY_CYCLE = ("edit", "already_correct", "edit", "negation") |
| MAX_WINDOW = 256 |
| MAX_COMMAND = 96 |
|
|
|
|
| def sha256(path: Path) -> str: |
| return hashlib.sha256(path.read_bytes()).hexdigest() |
|
|
|
|
| def input_key(command: str, binary: bytes) -> str: |
| """Length-delimited identity of the actual model input, not its label.""" |
| raw = command.encode("utf-8") |
| return hashlib.sha256(len(raw).to_bytes(8, "little") + raw + binary).hexdigest() |
|
|
|
|
| def signed32(value: int) -> int: |
| return (value + 2**31) % 2**32 - 2**31 |
|
|
|
|
| def expected(operation: str, inputs: list[int]) -> int: |
| left, right = inputs |
| if operation == "add": |
| return signed32(left + right) |
| if operation == "sub": |
| return signed32(left - right) |
| raise ValueError("unsupported arithmetic operation") |
|
|
|
|
| def make_row(rng: random.Random, index: int) -> dict: |
| family = FAMILY_CYCLE[index % len(FAMILY_CYCLE)] |
| source, requested = rng.choice((("add", "sub"), ("sub", "add"))) |
| actual = requested if family == "already_correct" else source |
| desired = actual if family == "negation" else requested |
| r1, r2, out = (f"x{x}" for x in rng.sample(TRAIN_REGS, 3)) |
| prefix = rng.randint(0, 16) |
| words = [rv32.addi("x0", "x0", rng.randint(-2048, 2047)) for _ in range(prefix)] |
| words.extend((rv32.lw(r1, 0, "x0"), rv32.lw(r2, 4, "x0"))) |
| location = len(words) |
| words.append(getattr(rv32, actual)(out, r1, r2)) |
| words.extend(rv32.addi("x0", "x0", rng.randint(-2048, 2047)) for _ in range(rng.randint(0, 7))) |
| words.extend((rv32.addi("x10", out, 0), rv32.EBREAK)) |
| target_words = list(words) |
| target_words[location] = getattr(rv32, desired)(out, r1, r2) |
| pack = lambda ws: b"".join(w.to_bytes(4, "little") for w in ws) |
| binary, target = pack(words), pack(target_words) |
| different = [i for i, (old, new) in enumerate(zip(binary, target)) if old != new] |
| if len(different) != int(family == "edit"): |
| raise AssertionError("generator did not produce its declared edit count") |
| pointer = different[0] if different else None |
| template = rng.choice(NEGATION_TEMPLATES if family == "negation" else TRAIN_TEMPLATES) |
| command = template.format(a=source, b=requested) |
| if len(binary) > MAX_WINDOW or len(command.encode()) > MAX_COMMAND: |
| raise AssertionError("public training input exceeds declared bounds") |
| return { |
| "id": index, |
| "family": family, |
| "command": command, |
| "binary_hex": binary.hex(), |
| "target_hex": target.hex(), |
| "pointer": pointer, |
| "replacement": target[pointer] if pointer is not None else 0, |
| "actual_operation": actual, |
| "desired_operation": desired, |
| "source_operation": source, |
| "requested_operation": requested, |
| "registers": [r1, r2, out], |
| "prefix_words": prefix, |
| "template": template, |
| "input_key": input_key(command, binary), |
| } |
|
|
|
|
| def check_row(row: dict, semantic: bool = False) -> None: |
| binary, target = bytes.fromhex(row["binary_hex"]), bytes.fromhex(row["target_hex"]) |
| if not 0 < len(binary) <= MAX_WINDOW or len(target) != len(binary): |
| raise ValueError("invalid training binary size") |
| if not 0 < len(row["command"].encode()) <= MAX_COMMAND: |
| raise ValueError("invalid training command size") |
| if row["input_key"] != input_key(row["command"], binary): |
| raise ValueError("training input key mismatch") |
| differences = [i for i, (old, new) in enumerate(zip(binary, target)) if old != new] |
| pointer = row["pointer"] |
| if pointer is None: |
| if differences or row["family"] == "edit" or row["replacement"] != 0: |
| raise ValueError("invalid NONE supervision") |
| elif differences != [pointer] or row["replacement"] != target[pointer] or row["family"] != "edit": |
| raise ValueError("invalid byte supervision") |
| if semantic: |
| edges = (-2**31, -1, 0, 1, 2**31 - 1) |
| for left in edges: |
| for right in edges: |
| inputs = [left, right] |
| if rv32.simulate(binary, inputs) != expected(row["actual_operation"], inputs): |
| raise ValueError("source execution disagrees with independent arithmetic") |
| if rv32.simulate(target, inputs) != expected(row["desired_operation"], inputs): |
| raise ValueError("target execution disagrees with independent arithmetic") |
|
|
|
|
| def freeze(output: Path, seed: int, count: int) -> dict: |
| if not 0 < count <= 200_000: |
| raise ValueError("training count must be between 1 and 200000") |
| output.mkdir(parents=True, exist_ok=True) |
| if any(output.iterdir()): |
| raise ValueError("refusing to overwrite a nonempty dataset directory") |
| rng = random.Random(seed) |
| rows, identities, keys = [], [], set() |
| attempts = 0 |
| for index in range(count): |
| while True: |
| row = make_row(rng, index) |
| attempts += 1 |
| if row["input_key"] not in keys: |
| break |
| check_row(row, semantic=index < 32) |
| keys.add(row["input_key"]) |
| rows.append(row) |
| identities.append({"id": index, "input_key": row["input_key"], "command": row["command"], |
| "binary_sha256": hashlib.sha256(bytes.fromhex(row["binary_hex"])).hexdigest()}) |
| dataset = output / "train.jsonl" |
| dataset.write_text("".join(json.dumps(r, sort_keys=True) + "\n" for r in rows)) |
| key_file = output / "seen-inputs.jsonl" |
| key_file.write_text("".join(json.dumps(r, sort_keys=True) + "\n" for r in identities)) |
| meta = { |
| "seed": seed, |
| "count": count, |
| "attempts": attempts, |
| "families": dict(Counter(r["family"] for r in rows)), |
| "operations": dict(Counter(r["desired_operation"] for r in rows)), |
| "train_registers": list(TRAIN_REGS), |
| "edit_templates": list(TRAIN_TEMPLATES), |
| "negation_templates": list(NEGATION_TEMPLATES), |
| "max_window": MAX_WINDOW, |
| "max_command": MAX_COMMAND, |
| "observed_max_binary_bytes": max(len(r["binary_hex"]) // 2 for r in rows), |
| "observed_max_command_bytes": max(len(r["command"].encode()) for r in rows), |
| "dataset_sha256": sha256(dataset), |
| "input_keys_sha256": sha256(key_file), |
| "generator_sha256": sha256(Path(__file__)), |
| "supervision": "assembler-supplied pointer/NONE and absolute replacement byte; not execution-only RL", |
| "ordering": "stored rows are shared by every model and seed; the training protocol fixes their execution order", |
| "validation": "all exact labels checked; first 32 sources and targets each executed on 25 independent arithmetic inputs", |
| } |
| (output / "dataset.json").write_text(json.dumps(meta, indent=2)) |
| print(json.dumps(meta), flush=True) |
| return meta |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--seed", type=int, required=True) |
| parser.add_argument("--count", type=int, required=True) |
| args = parser.parse_args() |
| freeze(args.output, args.seed, args.count) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
