Release prep: ensemble-aware CLI/server, clean-install validation, compact Run locally section and local inference doc
5f15b97 verified | """Command line interface for Pluto.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| from .loader import load_engine | |
| def _resolve_questions(value: str) -> dict: | |
| if os.path.exists(value): | |
| with open(value, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| return json.loads(value) | |
| def main(argv: list[str] | None = None) -> None: | |
| ap = argparse.ArgumentParser(prog="pluto", description="Pluto decision models") | |
| ap.add_argument("--model", default=os.environ.get("PLUTO_MODEL", "omida-ai/pluto-1")) | |
| ap.add_argument("--device", default=None, help="cpu, cuda, or mps; default is automatic") | |
| sub = ap.add_subparsers(dest="cmd", required=True) | |
| sub.add_parser("info", help="show checkpoint metadata") | |
| p = sub.add_parser("predict", help="answer typed questions about a state") | |
| p.add_argument("--state", help="state text") | |
| p.add_argument("--state-file", help="read state from a file") | |
| p.add_argument("--questions", required=True, help="questions JSON or path") | |
| p.add_argument("--max-len", type=int, default=None) | |
| s = sub.add_parser("serve", help="start the HTTP server") | |
| s.add_argument("--host", default="0.0.0.0") | |
| s.add_argument("--port", type=int, default=8000) | |
| args = ap.parse_args(argv) | |
| if args.cmd == "info": | |
| print(json.dumps(load_engine(args.model, device=args.device).info, indent=2)) | |
| return | |
| if args.cmd == "predict": | |
| if args.state_file: | |
| with open(args.state_file, "r", encoding="utf-8") as f: | |
| state = f.read() | |
| elif args.state is not None: | |
| state = args.state | |
| else: | |
| state = sys.stdin.read() | |
| engine = load_engine(args.model, device=args.device) | |
| result = engine.predict(state, _resolve_questions(args.questions), max_len=args.max_len) | |
| print(json.dumps(result, indent=2)) | |
| return | |
| if args.cmd == "serve": | |
| import uvicorn | |
| from .serve import create_app | |
| uvicorn.run(create_app(args.model, device=args.device), host=args.host, port=args.port) | |
| if __name__ == "__main__": | |
| main() | |