| """`zenith` command line: simulate, ask and eval.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| from zenith.engines.base import EngineError, SpecError |
|
|
|
|
| def _fail(msg: str) -> None: |
| print(f"error: {msg}", file=sys.stderr) |
| sys.exit(2) |
|
|
|
|
| def _simulate(args) -> None: |
| from zenith.engines.registry import simulate |
|
|
| try: |
| result = simulate(json.loads(args.spec)) |
| except json.JSONDecodeError as e: |
| _fail(f"spec is not valid JSON: {e}") |
| except (SpecError, EngineError) as e: |
| _fail(str(e)) |
| out = result.to_dict() |
| out.pop("series") |
| print(json.dumps(out, indent=2, ensure_ascii=False)) |
| if args.plot: |
| from zenith.engines.plot import save_plot |
|
|
| save_plot(result, Path(args.plot)) |
|
|
|
|
| def _ask(args) -> None: |
| from zenith.serve.pipeline import Zenith |
|
|
| zenith = Zenith.load(Path(args.ckpt), Path(args.tokenizer), Path(args.config), device=resolve_device(args.device)) |
| r = zenith.ask(args.question, temperature=args.temperature, max_answer_tokens=args.max_answer_tokens) |
| print(json.dumps({"spec": r.spec, "error": r.error, "truncated": r.truncated}, indent=2, ensure_ascii=False)) |
| if r.answer is not None: |
| print(r.answer) |
| if args.plot and r.result and r.result["series"]: |
| from zenith.engines.base import Result |
| from zenith.engines.plot import save_plot |
|
|
| save_plot(Result(**r.result), Path(args.plot)) |
| if r.error: |
| sys.exit(1) |
|
|
|
|
| def resolve_device(device: str) -> str: |
| """`auto` means the GPU when torch can see one, else the CPU.""" |
| if device != "auto": |
| return device |
| import torch |
|
|
| return "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
| def _eval(args) -> None: |
| from zenith.evals.run_eval import run_eval, run_showcase |
| from zenith.sft.generate import Example |
|
|
| if args.baseline == "rules": |
| from zenith.evals.baselines import RuleRouter |
|
|
| predictor = RuleRouter() |
| else: |
| if not args.ckpt: |
| _fail("--ckpt is required unless --baseline is given") |
| from zenith.serve.pipeline import Zenith |
|
|
| predictor = Zenith.load(Path(args.ckpt), Path(args.tokenizer), Path(args.config), |
| device=resolve_device(args.device)) |
| with open(args.data) as f: |
| examples = [Example.from_json(l) for l in f if l.strip()] |
| if args.limit: |
| step = max(1, len(examples) // args.limit) |
| examples = examples[::step][: args.limit] |
| report = run_eval(predictor, examples, Path(args.out)) |
| if args.showcase: |
| run_showcase(predictor, Path(args.out) / "showcase.md") |
| print((Path(args.out) / "report.md").read_text()) |
| print(json.dumps(report["overall"], indent=2)) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(prog="zenith", description="FarZenith astronomy model") |
| sub = parser.add_subparsers(dest="command", required=True) |
| p = sub.add_parser("simulate", help="run a physics engine on a JSON spec") |
| p.add_argument("spec", help='e.g. \'{"sim":"orbit","op":"kepler","star_mass_msun":2}\'') |
| p.add_argument("--plot", help="save the result's plots to this PNG path") |
| p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda") |
| p.set_defaults(func=_simulate) |
|
|
| p = sub.add_parser("ask", help="ask the model a question") |
| p.add_argument("question") |
| p.add_argument("--ckpt", required=True) |
| p.add_argument("--tokenizer", default="zenith/tokenizer/zenith-16k.json") |
| p.add_argument("--config", default="configs/small.yaml") |
| p.add_argument("--temperature", type=float, default=0.0) |
| p.add_argument("--max-answer-tokens", type=int, default=256) |
| p.add_argument("--plot", help="save the result's plots to this PNG path") |
| p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda") |
| p.set_defaults(func=_ask) |
|
|
| p = sub.add_parser("eval", help="score a checkpoint or baseline on eval examples") |
| p.add_argument("--ckpt") |
| p.add_argument("--baseline", choices=["rules"]) |
| p.add_argument("--data", required=True, help="eval.jsonl from zenith.sft.generate") |
| p.add_argument("--out", required=True) |
| p.add_argument("--tokenizer", default="zenith/tokenizer/zenith-16k.json") |
| p.add_argument("--config", default="configs/small.yaml") |
| p.add_argument("--limit", type=int) |
| p.add_argument("--showcase", action="store_true") |
| p.add_argument("--device", default="auto", help="auto (GPU if available), cpu or cuda") |
| p.set_defaults(func=_eval) |
| return parser |
|
|
|
|
| def main(argv=None) -> None: |
| args = build_parser().parse_args(argv) |
| args.func(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
