import argparse import asyncio import hashlib import importlib.metadata import json import time from pathlib import Path import tinker from tinker_cookbook import renderers from tinker_cookbook.completers import TinkerTokenCompleter from tinker_cookbook.recipes.math_rl import math_env from remote_stages import create from rx_gs1 import MODEL, RENDERER, ROOT, dataset, guard_grader, service, write_json def canonical(value): return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True, allow_nan=False) def digest(value): return hashlib.sha256(canonical(value).encode()).hexdigest() def summarize(rows): if not rows or len({r["item_id"] for r in rows}) != len(rows): raise ValueError("empty or duplicate evaluation items") if any(r["value"] not in (0, 1) for r in rows): raise ValueError("nonbinary accuracy") return dict(items=len(rows), correct=sum(r["value"] for r in rows), accuracy=sum(r["value"] for r in rows) / len(rows)) async def freeze(): _, ds = await dataset(1319, 1) assert len(ds.ds) == 1319 builders = list(ds.get_batch(0)) assert len(builders) == 1319 manifest = [] for i, builder in enumerate(builders): env = builder.env_thunk() prompt, stop = await env.initial_observation() manifest.append(dict(item_id=f"gsm8k-test-{i:04d}", question=env.problem, messages=env.convo_prefix + [dict(role="user", content=env.get_question())], prompt_tokens=prompt.to_ints(), reference=env.get_reference_answer())) manifest_text = "".join(canonical(row) + "\n" for row in manifest) versions = {name: importlib.metadata.version(name) for name in ("tinker", "tinker-cookbook", "datasets", "transformers", "tokenizers", "sympy", "numpy", "torch", "research-exchange")} files = [ROOT / "eval.py", ROOT / "rx_gs1.py"] files += [ROOT / "cookbook/tinker_cookbook" / name for name in ("recipes/math_rl/math_env.py", "recipes/math_rl/math_grading.py", "rl/problem_env.py", "completers.py")] files += sorted((ROOT / "cookbook/tinker_cookbook/renderers").glob("*.py")) code = dict(dependency_versions=versions, files={str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in files}) pin = json.loads((ROOT / "pins.json").read_text()) provenance = dict(code_sha256=digest(code), code_manifest=code, prompt_template=dict(renderer=RENDERER, tokenizer=MODEL, tokenizer_chat_template=ds.renderer.tokenizer.chat_template, conversation_prefix=math_env.MathEnv.standard_fewshot_prefix(), user_message="{gsm8k.question}" + math_env.MathEnv.question_suffix(), rendering="MathEnv.initial_observation -> renderer.build_generation_prompt; exact rendered token IDs in manifest"), decoding=dict(temperature=0, max_tokens=512, stop=stop, num_samples=1), parser=dict(description="Cookbook MathEnv.check_answer -> extract_boxed -> safe_grade(default sympy, timeout=1s) -> grade_answer: normalized/symbolic matching, NOT research_exchange numeric_exact. numeric_exact applies only to generation_audit, inapplicable here. Missing boxes/invalid outputs are incorrect; renderer.parse_response precedes grading. Accuracy excludes format reward.", cookbook_sha=pin.get("cookbook_sha", pin.get("cookbook_commit")), guard="math_grading._is_int(x) = math.isfinite(x) and original_is_int(x)", immutable_sources=code["files"]), items=dict(ids=[r["item_id"] for r in manifest], manifest_ref="evaluation-items.jsonl", manifest_sha256=hashlib.sha256(manifest_text.encode()).hexdigest(), dataset="openai/gsm8k", config="main", split="test", order="original dataset order")) evaluator_id = digest(provenance) frozen = dict(version=1, evaluator_id=evaluator_id, evaluators={evaluator_id: provenance}) path = ROOT / "evaluator.json" if path.exists() and json.loads(path.read_text()) != frozen: raise ValueError("evaluator already frozen with different provenance; do not overwrite") (ROOT / "evaluation-items.jsonl").write_text(manifest_text) write_json(path, frozen) return frozen, builders async def evaluate_base(): start = time.monotonic() guard_grader() frozen, builders = await freeze() evaluator_id = frozen["evaluator_id"] checkpoint = json.loads((ROOT / "init-save-response.json").read_text())["path"] persistent = ROOT / "base" transient = Path("/tmp/rx-gs1/base") persistent.mkdir(parents=True, exist_ok=True) transient.mkdir(parents=True, exist_ok=True) if (persistent / "eval.jsonl").exists(): raise ValueError("base evaluation already exists; refusing to overwrite") sc = service() status = "errored" timings = dict() try: tc = await create(sc) future = await tc.load_state_async(checkpoint) loaded = await future.result_async() future = await tc.save_weights_for_sampler_async(name="rx-base-eval") saved = await future.result_async() write_json(persistent / "server-responses.json", dict(load=loaded, sampler_save=saved)) sampler = await sc.create_sampling_client_async(model_path=saved.path) policy = TinkerTokenCompleter(sampler, max_tokens=512, temperature=0) timings["setup_seconds"] = time.monotonic() - start t = time.monotonic() rows = [] with (persistent / "eval.jsonl").open("w") as out, (transient / "eval.jsonl").open("w") as mirror: async def score(i, builder): env = builder.env_thunk() prompt, stop = await env.initial_observation() action = await policy(prompt, stop) message, _ = env.renderer.parse_response(action.tokens) prediction = renderers.get_text_content(message) row = dict(item_id=f"gsm8k-test-{i:04d}", arm="base", seed=0, evaluator_id=evaluator_id, metric="accuracy", prediction=prediction, reference=env.get_reference_answer(), value=int(env.check_answer(prediction)), checkpoint=checkpoint, sampler_checkpoint=saved.path) line = canonical(row) + "\n" out.write(line); out.flush() mirror.write(line); mirror.flush() rows.append(row) if len(rows) % 100 == 0: print(f"base scored {len(rows)}/1319", flush=True) await asyncio.gather(*(score(i, b) for i, b in enumerate(builders))) timings["scoring_seconds"] = time.monotonic() - t result = dict(summarize(rows), evaluator_id=evaluator_id, checkpoint=checkpoint, metric="accuracy", arm="base", seed=0, predictions_ref="base/eval.jsonl", timings=timings) write_json(ROOT / "evaluation-base.json", result) print(canonical(result), flush=True) status = "success" finally: await sc.close(status) timings["total_seconds"] = time.monotonic() - start write_json(persistent / "timings.json", timings) print("base evaluation session closed", flush=True) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("command", choices=["freeze", "base"]) args = parser.parse_args() if args.command == "freeze": print(asyncio.run(freeze())[0]["evaluator_id"]) else: asyncio.run(evaluate_base())