| """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 |
|
|
| 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: |
| 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") |
|
|
