Download scripts/run_benchmarks.py from ChatterjeeLab/DooABLe: direct link, hf CLI and curl.
- Browser
- Download file 4.47 kB
-
https://huggingface.co/ChatterjeeLab/DooABLe/resolve/main/scripts/run_benchmarks.py
- Command line
-
hf download hf://ChatterjeeLab/DooABLe/scripts/run_benchmarks.py
-
curl -L -o run_benchmarks.py https://huggingface.co/ChatterjeeLab/DooABLe/resolve/main/scripts/run_benchmarks.py
4.47 kB
| """Run the implemented exact and neural comparisons from a YAML config.""" | |
| import argparse, json, time, platform | |
| from pathlib import Path | |
| import numpy as np, pandas as pd, torch, yaml | |
| from dooable.graph import Graph, toy_graph, grid_graph, string_graph | |
| from dooable.exact import ( | |
| solve, | |
| uniform_policy, | |
| tilted_reference, | |
| backward_policy, | |
| forward_from_backward, | |
| sample, | |
| ) | |
| from dooable.learning import train | |
| from dooable.metrics import graph_metrics | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument("--config", required=True) | |
| a = p.parse_args() | |
| cfg = yaml.safe_load(Path(a.config).read_text()) | |
| out = Path(cfg["output"]) | |
| out.mkdir(parents=True, exist_ok=True) | |
| if cfg.get("graph"): | |
| g = Graph.load(cfg["graph"]) | |
| elif cfg.get("kind") == "grid": | |
| g = grid_graph(cfg.get("width", 5), cfg.get("budget", 6)) | |
| elif cfg.get("kind") == "strings": | |
| g = string_graph(cfg.get("length", 4), cfg.get("budget", 2)) | |
| else: | |
| g = toy_graph(cfg.get("multiplicity", 8)) | |
| rewards = ( | |
| json.loads(Path(cfg["rewards"]).read_text()) | |
| if cfg.get("rewards") | |
| else {y: 0.0 for y in g.terminals} | |
| ) | |
| temperature = cfg.get("temperature", 0.7) | |
| rows = [] | |
| g.save(out / "graph.json") | |
| for seed in cfg.get("seeds", [0, 1, 2, 3, 4]): | |
| for name in cfg["methods"]: | |
| start = time.perf_counter() | |
| if name == "exact": | |
| forward = solve(g, rewards, temperature).forward | |
| elif name == "uniform": | |
| forward = uniform_policy(g) | |
| elif name == "reference_tilt": | |
| forward = tilted_reference(g, rewards) | |
| elif name == "zero_cost": | |
| from dooable.ablations import zero_cost_policy | |
| forward = zero_cost_policy(g, rewards) | |
| elif name == "duplicate_endpoints": | |
| from dooable.ablations import duplicate_endpoint_policy | |
| forward = duplicate_endpoint_policy(g, rewards, temperature) | |
| elif name in ["dooable", "tb_uniform", "tb_exact", "unnormalized"]: | |
| backward = { | |
| "dooable": "learned", | |
| "tb_uniform": "uniform", | |
| "tb_exact": "exact", | |
| "unnormalized": "unnormalized", | |
| }[name] | |
| model, _ = train( | |
| g, | |
| rewards, | |
| temperature, | |
| steps=cfg.get("steps", 2000), | |
| batch_size=cfg.get("batch_size", 64), | |
| seed=seed, | |
| output=out / f"{name}_seed{seed}", | |
| backward=backward, | |
| ) | |
| forward = model.probabilities() | |
| else: | |
| raise ValueError(f"Unimplemented comparator {name}") | |
| row = { | |
| "method": name, | |
| "seed": seed, | |
| "seconds": time.perf_counter() - start, | |
| "nodes": len(g.nodes), | |
| "edges": len(g.edges), | |
| "outcomes": len(g.terminals), | |
| **graph_metrics(g, forward, rewards, temperature), | |
| } | |
| if g.metadata.get("kind") == "reaction": | |
| from dooable.chemistry import replay | |
| paths = sample(g, forward, cfg.get("samples", 1000), seed) | |
| row["replay_fraction"] = np.mean( | |
| [replay(r, g.metadata["budget"]) for r in paths] | |
| ) | |
| row["unique_outcomes"] = len({r["outcome"] for r in paths}) | |
| (out / f"{name}_seed{seed}_samples.jsonl").write_text( | |
| "".join(json.dumps(r) + "\n" for r in paths) | |
| ) | |
| rows.append(row) | |
| pd.DataFrame(rows).to_csv(out / "metrics.csv", index=False) | |
| print(json.dumps(row), flush=True) | |
| df = pd.DataFrame(rows) | |
| numeric = [c for c in df.select_dtypes("number").columns if c != "seed"] | |
| df.groupby("method")[numeric].agg(["mean", "sem"]).to_csv(out / "summary.csv") | |
| (out / "run.json").write_text( | |
| json.dumps( | |
| { | |
| "config": cfg, | |
| "python": platform.python_version(), | |
| "torch": torch.__version__, | |
| "numpy": np.__version__, | |
| "platform": platform.platform(), | |
| }, | |
| indent=2, | |
| ) | |
| ) | |
| if __name__ == "__main__": | |
| main() | |