File size: 2,781 Bytes
995bd7d b95ffe7 995bd7d | 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 | """aubin-mini command line.
aubin-mini decide request.json # {"state": ..., "questions": {...}} → JSON result
aubin-mini serve --port 8009 # POST /decide with the same JSON, GET /health
"""
import argparse
import json
import sys
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from . import AubinMini, __version__
def main(argv=None):
ap = argparse.ArgumentParser(prog="aubin-mini")
ap.add_argument("--model", default="emrevrg/AUBIN-Mini", help="HF repo or local directory")
ap.add_argument("--variant", default="phone", choices=["phone", "nb4"])
ap.add_argument("--threads", type=int, default=None)
sub = ap.add_subparsers(dest="cmd", required=True)
d = sub.add_parser("decide"); d.add_argument("request", help="JSON file, or - for stdin")
s = sub.add_parser("serve"); s.add_argument("--host", default="127.0.0.1"); s.add_argument("--port", type=int, default=8009)
a = ap.parse_args(argv)
m = AubinMini.from_pretrained(a.model, a.variant, threads=a.threads)
if a.cmd == "decide":
req = json.load(sys.stdin if a.request == "-" else open(a.request, encoding="utf-8"))
print(json.dumps(m.decide(req.get("state", ""), req["questions"]), ensure_ascii=False, indent=2))
return
class H(BaseHTTPRequestHandler):
def _send(self, code, obj):
b = json.dumps(obj, ensure_ascii=False).encode("utf-8")
self.send_response(code); self.send_header("Content-Type", "application/json; charset=utf-8")
self.send_header("Content-Length", str(len(b))); self.end_headers(); self.wfile.write(b)
def do_GET(self):
if self.path == "/health":
self._send(200, {"ok": True, "model": "AUBIN Mini", "version": __version__, "variant": m.variant})
else:
self._send(404, {"error": "not found"})
def do_POST(self):
if self.path != "/decide":
return self._send(404, {"error": "not found"})
try:
req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
self._send(200, m.decide(req.get("state", ""), req["questions"]))
except KeyError as e:
self._send(400, {"error": f"missing field: {e.args[0]}"})
except (ValueError, TypeError, AttributeError) as e:
self._send(400, {"error": str(e)})
def log_message(self, *args):
pass
print(f"AUBIN Mini {__version__} ({m.variant}) on http://{a.host}:{a.port} POST /decide", flush=True)
ThreadingHTTPServer((a.host, a.port), H).serve_forever()
if __name__ == "__main__":
main()
|