add serve/system1_server.py

serve/system1_server.py

@@ -0,0 +1,198 @@
+"""Local HTTP server speaking TypeSafe's System One contract (POST /v1/systemone) on top of a System1 backend.
+Run: python system1_server.py --model <path> [--adapter <safetensors|dir>] [--port 8020] [--model-name system1-9b]
+Probabilities come from the same path as scripts/evaluate_system1.py predict (render_messages -> backend.score ->
+calibrated_probabilities at T=1). Confidence formulas mirror the Kev reference (Apache-2.0, kev/api.py: choice_confidence,
+score_confidence, round_prob), themselves mirroring TypeSafe system-one-adapter 0.2.1 confidence_metrics."""
+import argparse
+import json
+import threading
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from kev_convert import kev_record_to_rows
+from system1_client import (
+    WIDE_CHOICE_KEYS, MAX_PAYLOAD_CHARS, OverlengthError, calibrated_probabilities, render_messages, request_from_row)
+
+MAX_BODY_BYTES = 4 * MAX_PAYLOAD_CHARS  # generous: JSON escaping and multi-question requests; rows are re-bounded per question
+
+
+class ContractError(ValueError):
+    pass
+
+
+def round_prob(x):
+    return round(float(x), 4)
+
+
+def _normalize(p):
+    t = sum(p)
+    return [1 / len(p)] * len(p) if t == 0 else [x / t for x in p]
+
+
+def choice_confidence(p):
+    K = len(p)
+    return 1.0 if K == 1 else (max(_normalize(p)) - 1 / K) / (1 - 1 / K)
+
+
+def score_confidence(p):
+    L = len(p)
+    if L == 1:
+        return 1.0
+    p = _normalize(p)
+    mode = max(range(L), key=p.__getitem__)
+    D = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
+    return max(0.0, 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / D)
+
+
+def request_to_rows(req):
+    """Validate a System One request -> ({qid: row}, {qid: candidate names}). Rows come from kev_convert with a dummy label
+    (never rendered), so option text equals what the evaluator and Jev conversion produce."""
+    if not isinstance(req, dict):
+        raise ContractError("request: expected object")
+    if "state" not in req or not isinstance(req["state"], (str, dict, list)):
+        raise ContractError("state: required, string, object or array")
+    qs = req.get("questions")
+    if not isinstance(qs, dict) or not qs:
+        raise ContractError("questions: expected nonempty object")
+    dummy, rows, names = {}, {}, {}
+    for qid, q in qs.items():
+        if not isinstance(q, dict):
+            raise ContractError(f"{qid}: expected object")
+        kind, crit = q.get("type"), q.get("criteria")
+        if kind == "choice":
+            if not isinstance(crit, (dict, list)) or not 2 <= len(crit) <= len(WIDE_CHOICE_KEYS):
+                raise ContractError(f"{qid}: choice requires 2-{len(WIDE_CHOICE_KEYS)} options per choice")
+            label = next(iter(crit))
+        elif kind == "score":
+            if not isinstance(crit, list):
+                raise ContractError(f"{qid}: score criteria must be an ordered list")
+            label = 0
+        elif kind == "noul":
+            if crit is not None and not isinstance(crit, dict):
+                raise ContractError(f"{qid}: noul criteria must be an object")
+            label = True
+        else:
+            raise ContractError(f"{qid}: unknown question type {kind!r}")
+        dummy[qid] = {"type": kind, "instructions": q.get("instructions"), "criteria": crit, "label": label}
+    try:
+        for row in kev_record_to_rows({"state": req["state"], "questions": dummy}, "serve", "none", "serve", keys=WIDE_CHOICE_KEYS):
+            qid = row["meta"]["kev_question_id"]
+            rows[qid] = row
+            names[qid] = list(row["meta"].get("candidate_ids", {}).values())
+    except (ValueError, KeyError, TypeError) as e:
+        raise ContractError(str(e))
+    return rows, names
+
+
+def answer_for(row, names, probs):
+    kind = row["type"]
+    if kind == "noul":
+        return {"type": "noul", "noul": round_prob(probs["Y"])}
+    if kind == "choice":
+        p = [probs[k] for k in row["options"]]
+        return {"type": "choice", "choice": names[max(range(len(p)), key=p.__getitem__)], "confidence": round_prob(choice_confidence(p)),
+                "probabilities": {n: round_prob(v) for n, v in zip(names, p)}}
+    p = [probs[k] for k in row["options"]]
+    return {"type": "score", "score": round_prob(sum(i * v for i, v in enumerate(p))),
+            "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)}, "confidence": round_prob(score_confidence(p))}
+
+
+class Service:
+    def __init__(self, backend, model_name):
+        self.backend, self.model_name, self.lock = backend, model_name, threading.Lock()  # ponytail: one global lock, GPU is serial anyway
+
+    def _tokens(self, text_or_messages):
+        tok = getattr(self.backend, "tokenizer", None)
+        if tok is None:
+            return 0
+        if isinstance(text_or_messages, str):
+            return len(tok.encode(text_or_messages, add_special_tokens=False))
+        prompt = tok.apply_chat_template(text_or_messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        return len(tok.encode(prompt, add_special_tokens=False))
+
+    def systemone(self, req):
+        start = time.perf_counter()
+        rows, names = request_to_rows(req)
+        answers, in_tokens = {}, 0
+        for qid, row in rows.items():
+            try:
+                messages = render_messages(request_from_row(row))
+            except OverlengthError as e:
+                raise ContractError(f"{qid}: payload too large, over the maximum context length: {e}")
+            except ValueError as e:
+                raise ContractError(f"{qid}: {e}")
+            with self.lock:
+                try:
+                    scores = self.backend.score(messages, list(row["options"]))
+                except OverlengthError as e:
+                    raise ContractError(f"{qid}: context too large, over the maximum context length: {e}")
+                in_tokens += self._tokens(messages)
+            if not isinstance(scores, dict) or set(scores) != set(row["options"]):
+                raise RuntimeError(f"{qid}: backend returned incomplete scores")
+            answers[qid] = answer_for(row, names[qid], calibrated_probabilities(scores, 1.0))
+        return {"model": self.model_name, "answers": answers,
+                "usage": {"input_tokens": in_tokens, "output_tokens": self._tokens(json.dumps(answers))},
+                "latency_ms": round((time.perf_counter() - start) * 1000, 1)}
+
+
+def make_server(backend, model_name="system1", host="127.0.0.1", port=8020):
+    svc = Service(backend, model_name)
+
+    class Handler(BaseHTTPRequestHandler):
+        def _send(self, code, body):
+            data = json.dumps(body).encode()
+            self.send_response(code)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+
+        def do_GET(self):
+            if self.path == "/health":
+                self._send(200, {"status": "ok"})
+            elif self.path == "/v1/models":
+                self._send(200, {"object": "list", "data": [{"id": svc.model_name, "object": "model", "owned_by": "local"}]})
+            else:
+                self._send(404, {"error": "not found"})
+
+        def do_POST(self):
+            if self.path != "/v1/systemone":
+                return self._send(404, {"error": "not found"})
+            n = int(self.headers.get("Content-Length") or 0)
+            if n > MAX_BODY_BYTES:
+                return self._send(422, {"error": f"payload too large: {n} bytes"})
+            try:
+                req = json.loads(self.rfile.read(n))
+            except (ValueError, UnicodeDecodeError) as e:
+                return self._send(400, {"error": f"malformed JSON: {e}"})
+            try:
+                self._send(200, svc.systemone(req))
+            except ContractError as e:
+                self._send(422, {"error": str(e)})
+            except Exception as e:
+                self._send(500, {"error": f"{type(e).__name__}: {e}"})
+
+        def log_message(self, *a):
+            pass
+
+    return ThreadingHTTPServer((host, port), Handler)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--adapter")
+    ap.add_argument("--port", type=int, default=8020)
+    ap.add_argument("--host", default="127.0.0.1")
+    ap.add_argument("--model-name", default="system1-9b")
+    ap.add_argument("--backend", choices=("mlx", "torch"), default="mlx", help="torch: CUDA via transformers")
+    a = ap.parse_args()
+    import system1_client
+    backend = system1_client.TorchBackend if a.backend == "torch" else system1_client.MlxBackend
+    server = make_server(backend(a.model, adapter_path=a.adapter), a.model_name, a.host, a.port)
+    print(f"serving {a.model_name} on http://{a.host}:{a.port}", flush=True)
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()