| """Inference server for node12 behind ``api.askdol.com`` at ``/naver`` (classifier-plan §6.2 option (a), §6.4). |
| |
| Run as ``python -m training.serve --package DIR --base DIR [--host 127.0.0.1 --port 8012]``. ``--package`` is a |
| pinned snapshot of the private HF repo (``export_hf`` layout: adapter/, heads.safetensors, taxonomy.json, |
| temperatures.json, config.json, MODEL_REVISION); ``--base`` is the local Gemma 4 12B checkpoint it was trained on. |
| |
| Endpoints (Bearer key from ``NAVER_PSY_CLASSIFIER_API_KEY``, compared in constant time, never logged): |
| - ``GET /naver/health`` -> repo, revision, model_revision, taxonomy (canonical), temperatures, fields, max_batch; |
| - ``POST /naver/predict`` ``{"inputs": [{"id", "segments": [{"kind", "text"}]}]}`` -> per input raw logits, |
| calibrated gated probabilities and labels, ``n_tokens`` and ``truncated``. Post text is never echoed or logged. |
| |
| The torch-free core ``Service`` takes a ``predict_fn(encoded ids) -> list of {field: logits}``, so it is tested |
| without a model; ``load_service`` builds the real one. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hmac |
| import json |
| import logging |
| import os |
| import re |
| from pathlib import Path |
| from typing import Any, Callable, Mapping |
|
|
| from .calibrate import outputs |
| from .formats import encode |
| from .taxonomy import Taxonomy |
|
|
| log = logging.getLogger("training.serve") |
| API_KEY_ENV = "NAVER_PSY_CLASSIFIER_API_KEY" |
| REVISION_ENV = "NAVER_PSY_MODEL_REVISION" |
| RENDER_VERSION = "nsb-input-v1" |
| PREFIX = "/naver" |
| REVISION_RE = re.compile(r"^[^@\s]+@[0-9a-f]{40}$") |
| MAX_SEGMENTS = 16 |
| MAX_TEXT = 20_000 |
|
|
| PredictFn = Callable[[list[list[int]]], list[dict[str, list[float]]]] |
|
|
|
|
| class RequestError(ValueError): |
| """A malformed request (HTTP 400/413); the message never contains post text.""" |
|
|
| def __init__(self, message: str, status: int = 400) -> None: |
| super().__init__(message) |
| self.status = status |
|
|
|
|
| class Service: |
| """Encode, predict, calibrate and gate one batch (torch-free).""" |
|
|
| def __init__(self, predict_fn: PredictFn, tokenize: Callable[[str], list[int]], bos: int | None, tax: Taxonomy, |
| temps: Mapping[str, float], model_revision: str, repo: str | None = None, |
| revision: str | None = None, max_len: int = 2048, max_batch: int = 64, |
| render_version: str = RENDER_VERSION) -> None: |
| if not REVISION_RE.match(model_revision): |
| raise ValueError("model_revision must be '<repo>@<40-hex commit>'") |
| r, c = model_revision.split("@") |
| if (repo and repo != r) or (revision and revision != c): |
| raise ValueError("repo/revision differ from model_revision") |
| if set(temps) != set(tax.names): |
| raise ValueError("temperatures do not cover the taxonomy fields") |
| self.predict_fn, self.tokenize, self.bos, self.tax = predict_fn, tokenize, bos, tax |
| self.temps = {f: float(temps[f]) for f in tax.names} |
| self.model_revision, self.repo, self.revision = model_revision, r, c |
| self.max_len, self.max_batch, self.render_version = max_len, max_batch, render_version |
|
|
| def health(self) -> dict[str, Any]: |
| return {"status": "ok", "repo": self.repo, "revision": self.revision, "model_revision": self.model_revision, |
| "taxonomy_version": self.tax.version, "render_version": self.render_version, |
| "taxonomy": self.tax.to_json(), "temperatures": self.temps, "fields": self.tax.names, |
| "max_batch": self.max_batch, "max_len": self.max_len} |
|
|
| def _validate(self, inputs: Any) -> list[tuple[str, list[dict]]]: |
| if not isinstance(inputs, list) or not inputs: |
| raise RequestError("inputs must be a non-empty list") |
| if len(inputs) > self.max_batch: |
| raise RequestError(f"batch of {len(inputs)} exceeds max_batch {self.max_batch}", 413) |
| out, seen = [], set() |
| for i, item in enumerate(inputs): |
| if not isinstance(item, Mapping) or not isinstance(item.get("id"), (str, int)): |
| raise RequestError(f"input {i}: needs an id") |
| rid = str(item["id"]) |
| if rid in seen: |
| raise RequestError(f"input {i}: duplicate id") |
| seen.add(rid) |
| segs = item.get("segments") |
| if not isinstance(segs, list) or not segs or len(segs) > MAX_SEGMENTS: |
| raise RequestError(f"input {i}: segments must be a list of 1-{MAX_SEGMENTS}") |
| clean = [] |
| for s in segs: |
| if (not isinstance(s, Mapping) or not isinstance(s.get("kind"), str) |
| or not isinstance(s.get("text"), str) or len(s["text"]) > MAX_TEXT): |
| raise RequestError(f"input {i}: each segment needs a string kind and text (<= {MAX_TEXT} chars)") |
| clean.append({"kind": s["kind"], "text": s["text"]}) |
| out.append((rid, clean)) |
| return out |
|
|
| def predict(self, inputs: Any) -> dict[str, Any]: |
| items = self._validate(inputs) |
| encoded, infos = [], [] |
| for i, (_, segs) in enumerate(items): |
| try: |
| ids, info = encode(segs, self.tokenize, self.max_len, self.bos) |
| except ValueError: |
| raise RequestError(f"input {i}: cannot be encoded within max_len") from None |
| encoded.append(ids) |
| infos.append(info) |
| logits = self.predict_fn(encoded) |
| if len(logits) != len(items): |
| raise RuntimeError("predict_fn returned a different number of rows") |
| results = [] |
| for (rid, _), lg, info in zip(items, logits, infos): |
| row = {f: [float(x) for x in lg[f]] for f in self.tax.names} |
| out = outputs(row, self.tax, self.temps) |
| results.append({"id": rid, "logits": row, "probs": out["probs"], "labels": out["labels"], |
| "n_tokens": info["n_tokens"], "truncated": bool(info["cut"])}) |
| return {"model_revision": self.model_revision, "results": results} |
|
|
|
|
| def authorised(header: str | None, key: str | None) -> bool: |
| """Constant-time Bearer check; no key configured means every request is refused.""" |
| if not key or not header or not header.startswith("Bearer "): |
| return False |
| return hmac.compare_digest(header[len("Bearer "):].encode(), key.encode()) |
|
|
|
|
| def create_app(service: Service, api_key: str | None = None, prefix: str = PREFIX): |
| """FastAPI app around a Service (FastAPI is imported here so the core needs nothing).""" |
| from fastapi import FastAPI, Request |
| from fastapi.responses import JSONResponse |
|
|
| key = api_key if api_key is not None else os.environ.get(API_KEY_ENV) |
| app = FastAPI(title="naver-psy classifier", docs_url=None, redoc_url=None, openapi_url=None) |
|
|
| def denied() -> JSONResponse: |
| return JSONResponse({"error": "unauthorised"}, status_code=401) |
|
|
| def health(request): |
| if not authorised(request.headers.get("authorization"), key): |
| return denied() |
| return service.health() |
|
|
| async def predict(request): |
| if not authorised(request.headers.get("authorization"), key): |
| return denied() |
| try: |
| body = await request.json() |
| except Exception: |
| return JSONResponse({"error": "body must be JSON"}, status_code=400) |
| try: |
| return service.predict(body.get("inputs") if isinstance(body, dict) else None) |
| except RequestError as e: |
| return JSONResponse({"error": str(e)}, status_code=e.status) |
| except Exception as e: |
| log.error("predict failed: %s", type(e).__name__) |
| return JSONResponse({"error": "internal error"}, status_code=500) |
|
|
| |
| health.__annotations__["request"] = Request |
| predict.__annotations__["request"] = Request |
| app.get(prefix + "/health")(health) |
| app.post(prefix + "/predict")(predict) |
| return app |
|
|
|
|
| def read_revision(package: Path) -> str: |
| env = os.environ.get(REVISION_ENV) |
| if env: |
| return env.strip() |
| path = package / "MODEL_REVISION" |
| if not path.exists(): |
| raise SystemExit(f"no model revision: set {REVISION_ENV} or write {path}") |
| return path.read_text().strip() |
|
|
|
|
| def load_service(package_dir, base_dir, device: str = "cuda", max_batch: int = 64) -> Service: |
| """Build the torch model from an HF package snapshot and a local base checkpoint.""" |
| import torch |
|
|
| from .data import inference_batches |
| from .model import bos_id, build, load_tokenizer, predict, tokenizer_fn |
|
|
| package = Path(package_dir) |
| cfg = json.loads((package / "config.json").read_text()) |
| cfg["model"] = str(base_dir) |
| tax = Taxonomy.load(package / "taxonomy.json") |
| temps = json.loads((package / "temperatures.json").read_text()) |
| dev = torch.device(device) |
| model = build(cfg, tax.sizes(), dev, train=False, adapter_dir=package / "adapter", |
| heads_path=package / "heads.safetensors") |
| tok = load_tokenizer(str(base_dir)) |
| max_len = int(cfg.get("max_len", 2048)) |
| budget = int(cfg.get("serve_token_budget", 65536)) |
|
|
| def predict_fn(encoded: list[list[int]]) -> list[dict[str, list[float]]]: |
| batches = inference_batches([len(x) for x in encoded], budget) |
| out = predict(model, encoded, tok.pad_token_id, batches, dev, tax.names) |
| return [out[i] for i in range(len(encoded))] |
|
|
| return Service(predict_fn, tokenizer_fn(tok), bos_id(tok), tax, temps, read_revision(package), |
| max_len=max_len, max_batch=max_batch) |
|
|
|
|
| def main(argv=None) -> int: |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| p.add_argument("--package", required=True, help="Pinned HF package snapshot") |
| p.add_argument("--base", required=True, help="Local base checkpoint directory") |
| p.add_argument("--host", default="127.0.0.1") |
| p.add_argument("--port", type=int, default=8012) |
| p.add_argument("--device", default="cuda") |
| p.add_argument("--max-batch", type=int, default=64) |
| args = p.parse_args(argv) |
| if not os.environ.get(API_KEY_ENV): |
| raise SystemExit(f"{API_KEY_ENV} is not set; refusing to serve without auth") |
| import uvicorn |
|
|
| logging.basicConfig(level=logging.INFO) |
| service = load_service(args.package, args.base, args.device, args.max_batch) |
| log.info("serving %s", service.model_revision) |
| uvicorn.run(create_app(service), host=args.host, port=args.port, access_log=False) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
