File size: 2,225 Bytes
5f15b97 | 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 | """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()
|