File size: 2,735 Bytes
adc02fa | 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 | #!/usr/bin/env python
from __future__ import annotations
import argparse
import sys
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from dovla_cil.experiments.scaling import ( # noqa: E402
ScalingExperiment,
parse_k_values,
run_scaling_experiment,
)
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description="Run scaling-law experiments over intervention multiplicity K."
)
parser.add_argument("--backend", choices=["toy"], default="toy")
parser.add_argument("--tasks", default="builtins", help="'builtins' or TaskSpec JSON/JSONL path.")
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--total-records", type=int, default=4096)
parser.add_argument("--k-values", default="1,2,4,8,16,32")
parser.add_argument("--epochs", type=int, default=3)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--shard-size", type=int, default=1000)
parser.add_argument("--batch-groups", type=int, default=8)
parser.add_argument("--records-per-group", type=int, default=8)
parser.add_argument("--hidden-dim", type=int, default=256)
parser.add_argument("--lr", type=float, default=1e-3)
parser.add_argument("--device", default="auto")
parser.add_argument(
"--eval-num-tasks",
type=int,
default=20,
help="Number of toy CausalStress groups per K.",
)
parser.add_argument(
"--eval-k",
type=int,
default=None,
help="Override CausalStress K. Defaults to the current scaling K.",
)
args = parser.parse_args(argv)
config = ScalingExperiment(
backend=args.backend,
tasks=args.tasks,
output_dir=args.out,
total_records=args.total_records,
k_values=parse_k_values(args.k_values),
epochs=args.epochs,
seed=args.seed,
shard_size=args.shard_size,
batch_groups=args.batch_groups,
records_per_group=args.records_per_group,
hidden_dim=args.hidden_dim,
learning_rate=args.lr,
device=args.device,
eval_num_tasks=args.eval_num_tasks,
eval_k=args.eval_k,
)
print("planned runs:")
for run in config.planned_runs():
print(run)
summary = run_scaling_experiment(config)
print(f"wrote aggregate CSV to {summary['aggregate_csv']}")
print(f"wrote plots to {args.out}")
for metric, values in summary["regression"].items():
print(f"{metric}: beta_log_k={values['beta_log_k']:.6g}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|