File size: 4,681 Bytes
5216cd1 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """POST /v1/systemone server for FINAL-Bench/Darwin-27B-ZTC (System One request shape).
ZTC_MODEL=<snapshot path or FINAL-Bench/Darwin-27B-ZTC> PORT=8000 python ztc_server.py
Every question is answered with one forward pass of the backbone and the decision readout
(no generated tokens). Questions from concurrent requests are batched together on the GPU
(at most ZTC_BATCH questions and ZTC_TOKEN_BUDGET estimated tokens per batch). A question
longer than ZTC_MAX_LENGTH tokens is refused with HTTP 422 (never truncated).
"""
import asyncio
import json
import math
import os
import queue
import sys
import threading
from concurrent.futures import Future
from pathlib import Path
import torch
import uvicorn
from fastapi import FastAPI, HTTPException, Request
MODEL = os.getenv("ZTC_MODEL", "FINAL-Bench/Darwin-27B-ZTC")
MAX_LENGTH = int(os.getenv("ZTC_MAX_LENGTH", "32768"))
BATCH = int(os.getenv("ZTC_BATCH", "8"))
TOKEN_BUDGET = int(os.getenv("ZTC_TOKEN_BUDGET", "16000"))
path = Path(MODEL)
if not path.exists():
from huggingface_hub import snapshot_download
path = Path(snapshot_download(MODEL, revision=os.getenv("ZTC_REVISION")))
sys.path.insert(0, str(path))
from autojev.model import DecisionModel, answer # noqa: E402
model = DecisionModel(checkpoint=path).eval()
NAME = "darwin-27b-ztc"
Q: "queue.Queue[tuple[dict, Future]]" = queue.Queue()
def estimate(row):
return len(json.dumps([row["state"], row["question"]], ensure_ascii=False)) // 3 + 192
@torch.inference_mode()
def probabilities(rows):
batch = model.prepare(rows, max_length=MAX_LENGTH)
values = (model(batch) / model.temperature).softmax(-1).cpu().tolist()
return [v[:c] for v, c in zip(values, batch.counts)]
def run_single(row, future):
try:
future.set_result(probabilities([row])[0])
except ValueError as error:
future.set_exception(ValueError(str(error)))
except Exception as error: # out of memory or other runtime failure
torch.cuda.empty_cache()
future.set_exception(RuntimeError(str(error)[:300]))
def worker():
while True:
items = [Q.get()]
while len(items) < BATCH * 4:
try:
items.append(Q.get(timeout=0.005))
except queue.Empty:
break
items.sort(key=lambda item: estimate(item[0]))
i = 0
while i < len(items):
group = [items[i]]
tokens = estimate(items[i][0])
while i + len(group) < len(items) and len(group) < BATCH and tokens + estimate(items[i + len(group)][0]) <= TOKEN_BUDGET:
tokens += estimate(items[i + len(group)][0])
group.append(items[i + len(group)])
i += len(group)
try:
for (row, future), values in zip(group, probabilities([g[0] for g in group])):
future.set_result(values)
except Exception:
torch.cuda.empty_cache()
for row, future in group:
if not future.done():
run_single(row, future)
threading.Thread(target=worker, daemon=True).start()
app = FastAPI(title="Darwin-27B-ZTC System One server")
@app.get("/health")
def health():
return {"status": "ready", "model": NAME, "temperature": model.temperature}
@app.post("/v1/systemone")
async def system_one(request: Request):
body = await request.json()
questions = body.get("questions") or {}
if body.get("images"):
raise HTTPException(422, "This model is text-only; images are not supported.")
futures = {}
for key, question in questions.items():
if question.get("type") not in ("choice", "noul", "score"):
raise HTTPException(422, f"Unsupported question type: {question.get('type')}")
future = Future()
Q.put(({"state": body.get("state"), "question": question}, future))
futures[key] = future
answers = {}
for key, future in futures.items():
try:
values = await asyncio.wrap_future(future)
except ValueError as error:
raise HTTPException(422, str(error))
except RuntimeError as error:
raise HTTPException(500, str(error))
if not all(math.isfinite(v) and v >= 0 for v in values) or sum(values) <= 0:
raise HTTPException(500, 'Model produced an invalid probability vector; refusing to substitute a prediction.')
answers[key] = answer(questions[key], values)
return {"model": NAME, "answers": answers}
if __name__ == "__main__":
uvicorn.run(app, host=os.getenv("HOST", "127.0.0.1"), port=int(os.getenv("PORT", "8000")), log_level="warning")
|