Spaces:
Sleeping
Sleeping
File size: 3,637 Bytes
05a4ec0 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """
Random Toy-IR list-of-dicts for GRPO / Deliverable2 / runtime_core.
Schema matches `runtime_core.SAMPLE_PROGRAM` (op, args, dest, type) — the same
shape `Deliverable2_Formatter` and `CompilerOptimizationEnv` expect.
Note: `metahack1 (1).ipynb` uses a different TAC shape (CONST/src1/STORE).
Use that notebook's generators only if you add a separate converter; this
module is self-contained for training stack compatibility.
"""
from __future__ import annotations
import copy
import json
import random
from typing import List
# Hand-written seeds (same as Colab / train.py short list) — kept in sync for regression.
_BUILTIN_EXTRA: List[List[dict]] = [
[
{"op": "const", "dest": "a", "args": ["1"], "type": "int"},
{"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"},
{"op": "ret", "args": ["b"]},
],
[
{"op": "const", "dest": "x", "args": ["2"], "type": "int"},
{"op": "const", "dest": "y", "args": ["4"], "type": "int"},
{"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"},
{"op": "const", "dest": "k", "args": ["1"], "type": "int"},
{"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"},
{"op": "ret", "args": ["w"]},
],
]
def random_toy_ir_program(rng: random.Random) -> List[dict]:
"""
One valid program: consts v0.., then a chain of add/mul on existing names, then ret.
All ops use the mock-engine-friendly list schema.
"""
n_const = rng.randint(2, 5)
n_arith = rng.randint(1, 5)
progs: List[dict] = []
for i in range(n_const):
progs.append(
{
"op": "const",
"dest": f"v{i}",
"args": [str(rng.randint(0, 20))],
"type": "int",
}
)
available = [f"v{i}" for i in range(n_const)]
nxt = n_const
for _j in range(n_arith):
a = rng.choice(available)
b = rng.choice(available)
opn = rng.choice(["add", "mul"])
d = f"v{nxt}"
nxt += 1
progs.append({"op": opn, "dest": d, "args": [a, b], "type": "int"})
available.append(d)
progs.append({"op": "ret", "args": [available[-1]]})
return progs
def build_training_program_corpus(
n_total: int = 120,
seed: int = 42,
*,
include_builtins: bool = True,
) -> List[List[dict]]:
"""
Return `n_total` programs for GRPO. Optionally prepend SAMPLE_PROGRAM + 2 hand-written IRs
(when include_builtins), then fill with random_toy_ir_program, deduplicating by JSON key.
Typical range: set `n_total` between 50 and 200 in the notebook.
"""
if n_total < 1:
raise ValueError("n_total must be >= 1")
rng = random.Random(seed)
out: List[List[dict]] = []
seen: set[str] = set()
def _add(p: List[dict]) -> None:
k = json.dumps(p, sort_keys=True)
if k in seen:
return
seen.add(k)
out.append(copy.deepcopy(p))
if include_builtins:
from runtime_core import SAMPLE_PROGRAM
for p in (SAMPLE_PROGRAM, *_BUILTIN_EXTRA):
if len(out) >= n_total:
break
_add(p)
# Fill with random programs (dedupe by full JSON; allow dup if generator keeps colliding)
guard = 0
while len(out) < n_total:
guard += 1
if guard > 200_000:
out.append(random_toy_ir_program(rng))
continue
cand = random_toy_ir_program(rng)
k = json.dumps(cand, sort_keys=True)
if k in seen:
continue
seen.add(k)
out.append(cand)
return out
|