File size: 3,549 Bytes
3f2dde4 | 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 | from __future__ import annotations
import argparse
from dataclasses import asdict
from .controller import TrainerConfig, fit_controller
from .benchmarks import run_benchmark_suite
from .gui import launch_runtime_gui
from .smoke import run_tinygrad_gate_demo
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(prog="openpeer-trainer")
sub = parser.add_subparsers(dest="command", required=True)
demo = sub.add_parser("demo", help="run the tinygrad gate-controller smoke demo")
demo.add_argument("--steps", type=int, default=80)
demo.add_argument("--batch-size", type=int, default=64)
demo.add_argument("--seed", type=int, default=0)
demo.add_argument("--target-accuracy", type=float, default=0.99)
fit = sub.add_parser("fit", help="fit an ntkmirror controller on a frozen causal LM")
fit.add_argument("--model", default="OpenPeerAI/OpenPeerLLM")
fit.add_argument("--fallback-model", default="sshleifer/tiny-gpt2")
fit.add_argument("--train-jsonl")
fit.add_argument("--out", default="runs/openpeer_controller.pt")
fit.add_argument("--gates", type=int, default=512)
fit.add_argument("--steps", type=int, default=40)
fit.add_argument("--demo-mode", action="store_true")
bench = sub.add_parser("bench", help="run local gate-controller benchmarks and generate charts")
bench.add_argument("--steps", type=int, nargs="+", default=[20, 40, 80, 120])
bench.add_argument("--batch-size", type=int, default=64)
bench.add_argument("--seed", type=int, default=0)
bench.add_argument("--output-dir", default="artifacts/benchmarks")
bench.add_argument("--target-accuracy", type=float, default=0.99)
gui = sub.add_parser("gui", help="launch the runtime GUI")
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
if args.command == "demo":
result = run_tinygrad_gate_demo(steps=args.steps, batch_size=args.batch_size, seed=args.seed, target_accuracy=args.target_accuracy)
print("tinygrad gate demo")
print(f"initial_accuracy={result.initial_accuracy:.3f}")
print(f"final_accuracy={result.final_accuracy:.3f}")
print(f"final_loss={result.final_loss:.4f}")
print(f"target_accuracy={result.target_accuracy:.3f}")
print(f"reached_target={result.reached_target}")
print(f"trained_steps={result.trained_steps}")
print(f"learned_gate_sample={result.learned_gate_sample}")
return 0
if args.command == "bench":
result = run_benchmark_suite(
step_counts=args.steps,
batch_size=args.batch_size,
seed=args.seed,
output_dir=args.output_dir,
target_accuracy=args.target_accuracy,
)
print(f"benchmarks_csv={result.csv_path}")
print(f"charts={result.chart_paths}")
return 0
if args.command == "gui":
return launch_runtime_gui()
config = TrainerConfig(
model_name=args.model,
fallback_model_name=args.fallback_model,
train_jsonl=args.train_jsonl,
out_path=args.out,
gates=args.gates,
steps=args.steps,
demo_mode=args.demo_mode,
)
output_path = fit_controller(config)
print(f"saved_controller={output_path}")
print(f"config={asdict(config)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|