File size: 4,821 Bytes
0cf8154 | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """`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] # spread across ops, which are stored in order
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()
|