| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Structured decisions in front of a vLLM DiffusionGemma server. |
| 4 | |
| 5 | POST /v1/systemone takes Jev's request body: {"model", "state", "questions"}. |
| 6 | "questions" maps an id to {"type", "instructions", "criteria"}, where "type" |
| 7 | is "noul", "choice" or "score" and the criteria shape follows the type: |
| 8 | noul: optional {"true": ..., "false": ...} descriptions |
| 9 | choice: option name -> description or null |
| 10 | score: ordered list of levels |
| 11 | Answers take Jev's shapes, with this server's diagnostics alongside: |
| 12 | noul: {"noul": p} |
| 13 | choice: {"choice", "probabilities", "confidence"} |
| 14 | score: {"score", "legend", "probabilities", "confidence"} |
| 15 | The request body may also carry the schema keys "instructions", "samples", |
| 16 | "auto_max", "auto_threshold", "steps", "think", "ask", "chunk_rows", |
| 17 | "chunk_prompt" and "sequential" as extensions. Images go ahead of the |
| 18 | state, either as multipart/form-data with the JSON body in a part named |
| 19 | "request" and each image as a file part, or as an "images" array of data |
| 20 | URLs in the JSON body. |
| 21 | |
| 22 | POST /v1/chat/completions makes the same decision from an OpenAI-shaped |
| 23 | call. The system message is the schema JSON below and the user message is |
| 24 | the state. The reply's `content` is the JSON answer set. |
| 25 | |
| 26 | POST /v1/raw/chat/completions passes the body to vLLM's chat completions |
| 27 | unchanged, for plain generation through this port. |
| 28 | |
| 29 | With API_KEY set in the environment, every POST needs "Authorization: |
| 30 | Bearer <key>". --tls-port adds an HTTPS listener with a self-signed |
| 31 | certificate kept in --cert-dir, for clients that need a secure origin. |
| 32 | |
| 33 | Each answer is one distribution per question, from one denoise step over a |
| 34 | seeded canvas, averaged over a few noise draws. This server handles the |
| 35 | canvas, tokenizer, slot resolution, noise draws and averaging. |
| 36 | |
| 37 | Schema (system message): |
| 38 | {"questions": [ |
| 39 | {"id": "urgent", "type": "noul", "instructions": "..."}, |
| 40 | {"id": "bucket", "type": "choice", "instructions": "...", |
| 41 | "options": [{"name": "billing", "description": "..."}, ...]}, |
| 42 | {"id": "tone", "type": "score", "instructions": "...", |
| 43 | "levels": ["calm", "annoyed", "furious"]}], |
| 44 | "instructions": "optional context", |
| 45 | "samples": "auto" | N, "auto_threshold": 0.1, "auto_max": 4, |
| 46 | "steps": 1, "think": 0} |
| 47 | |
| 48 | A question may also declare: |
| 49 | "depends_on": [ids] answered after those, with their answers in |
| 50 | its prompt |
| 51 | "ask_if": {id: [answers]} asked only when that question's answer is |
| 52 | among them (a skipped answer is null) |
| 53 | "alone": true a read of its own |
| 54 | Questions run in stages by these dependencies. Each stage is one joint |
| 55 | read. Later stages continue the earlier answers, prefilled for a text |
| 56 | state and restated for an image. |
| 57 | |
| 58 | Up to ten questions answer as "id: label" lines. Past that the id runs |
| 59 | straight into the label, space separated, one row fewer per question. The |
| 60 | server splits a schema whose answer template does not fit the canvas into |
| 61 | chunks that run together, each with its own question list. "chunk_rows" |
| 62 | sets the rows per chunk, "ask" picks a subset of question ids for one |
| 63 | read, "sequential": true runs the chunks in order with the earlier answers |
| 64 | prefilled, and "chunk_prompt": "shared" lists every question in each |
| 65 | chunk's prompt. "think": N lets the model write up to N tokens in its |
| 66 | thought channel, as an ordinary generation, and the read then runs with |
| 67 | that thought in its prompt. With images the model writes the thought with |
| 68 | the image in view and the server seeds it into the canvas ahead of the |
| 69 | answer, so the canvas bounds it. The noise draws of a decision share one |
| 70 | thought. |
| 71 | |
| 72 | Serve the model with a canvas that holds the answer template, for example: |
| 73 | vllm serve google/diffusiongemma-26B-A4B-it \ |
| 74 | --diffusion-config '{"canvas_length": 64}' --max-logprobs 32 \ |
| 75 | --enable-prefix-caching |
| 76 | then run this in front of it: |
| 77 | python structured_server.py --upstream http://127.0.0.1:8000 \ |
| 78 | --tokenizer google/diffusiongemma-26B-A4B-it --canvas 64 --port 8011 |
| 79 | """ |
| 80 | |
| 81 | import argparse |
| 82 | import json |
| 83 | import math |
| 84 | import os |
| 85 | import random |
| 86 | import ssl |
| 87 | import subprocess |
| 88 | import threading |
| 89 | import time |
| 90 | import urllib.error |
| 91 | import urllib.request |
| 92 | from concurrent.futures import ThreadPoolExecutor |
| 93 | from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| 94 | |
| 95 | import pybase64 as base64 |
| 96 | from transformers import AutoTokenizer |
| 97 | |
| 98 | ARGS = None |
| 99 | TOK = None |
| 100 | API_KEY = os.environ.get( |
| 101 | "API_KEY", "" |
| 102 | ) # when set, POST routes need "Authorization: Bearer <key>" |
| 103 | CANVAS_LEN = 64 # the served canvas length. A request may be narrower. |
| 104 | CANVAS_STEP = 16 # request widths are multiples of this |
| 105 | VOCAB = 262144 |
| 106 | TURN_CLOSE = 106 |
| 107 | PAD = 0 |
| 108 | TOPK = 20 |
| 109 | MAX_QUESTIONS = 64 # per request |
| 110 | MAX_SAMPLES = 32 # reads per question, fixed or auto |
| 111 | MAX_PARALLEL = 16 # question groups read at once |
| 112 | # the empty thought block the chat template leaves to the model |
| 113 | SCAFFOLD_TEXT = "<|channel>thought\n<channel|>" |
| 114 | SCAFFOLD = None |
| 115 | THOUGHT_OPEN = None |
| 116 | THOUGHT_CLOSE = None |
| 117 | |
| 118 | |
| 119 | # ---------------------------------------------------------------------------- |
| 120 | # Schema |
| 121 | # ---------------------------------------------------------------------------- |
| 122 | |
| 123 | |
| 124 | class SchemaError(ValueError): |
| 125 | pass |
| 126 | |
| 127 | |
| 128 | def parse_schema(value): |
| 129 | if ( |
| 130 | not isinstance(value, dict) |
| 131 | or not isinstance(value.get("questions"), list) |
| 132 | or not value["questions"] |
| 133 | ): |
| 134 | raise SchemaError("schema: needs a non-empty questions array") |
| 135 | if len(value["questions"]) > MAX_QUESTIONS: |
| 136 | raise SchemaError(f"schema: at most {MAX_QUESTIONS} questions") |
| 137 | qs = [] |
| 138 | seen = set() |
| 139 | for q in value["questions"]: |
| 140 | qid = str(q.get("id", "")).strip() |
| 141 | if not qid or ":" in qid or "\n" in qid: |
| 142 | raise SchemaError( |
| 143 | f"question id {qid!r} must be non-empty, no ':' or newline" |
| 144 | ) |
| 145 | if qid in seen: |
| 146 | raise SchemaError(f"duplicate question id {qid!r}") |
| 147 | seen.add(qid) |
| 148 | kind = q.get("type") |
| 149 | if kind in ("noul", "bool", "boolean"): |
| 150 | kind = "noul" |
| 151 | crit = q.get("criteria") or {} |
| 152 | choices = [("yes", crit.get("true")), ("no", crit.get("false"))] |
| 153 | labels = ["yes", "no"] |
| 154 | elif kind == "choice": |
| 155 | opts = q.get("options") or [] |
| 156 | choices = [ |
| 157 | (o["name"], o.get("description")) |
| 158 | if isinstance(o, dict) |
| 159 | else (str(o), None) |
| 160 | for o in opts |
| 161 | ] |
| 162 | labels = [chr(ord("A") + i) for i in range(len(choices))] |
| 163 | elif kind == "score": |
| 164 | choices = [(str(level), None) for level in (q.get("levels") or [])] |
| 165 | labels = ( |
| 166 | [str(i + 1) for i in range(len(choices))] |
| 167 | if len(choices) <= 9 |
| 168 | else [chr(ord("A") + i) for i in range(len(choices))] |
| 169 | ) |
| 170 | else: |
| 171 | raise SchemaError(f"question {qid!r}: unknown type {kind!r}") |
| 172 | if len(choices) < 2: |
| 173 | raise SchemaError(f"question {qid!r}: needs at least two alternatives") |
| 174 | if len(choices) > 26: |
| 175 | raise SchemaError(f"question {qid!r}: at most 26 alternatives") |
| 176 | deps = q.get("depends_on") or [] |
| 177 | ask_if = q.get("ask_if") or {} |
| 178 | if not isinstance(deps, list) or not all(isinstance(d, str) for d in deps): |
| 179 | raise SchemaError( |
| 180 | f"question {qid!r}: depends_on must be a list of question ids" |
| 181 | ) |
| 182 | if not isinstance(ask_if, dict) or not all( |
| 183 | isinstance(v, list) and v for v in ask_if.values() |
| 184 | ): |
| 185 | raise SchemaError( |
| 186 | f"question {qid!r}: ask_if must map a question id to a " |
| 187 | "non-empty list of its answers" |
| 188 | ) |
| 189 | qs.append( |
| 190 | { |
| 191 | "id": qid, |
| 192 | "type": kind, |
| 193 | "instructions": str(q.get("instructions", "")), |
| 194 | "choices": choices, |
| 195 | "labels": labels, |
| 196 | "depends_on": list(dict.fromkeys(list(deps) + list(ask_if))), |
| 197 | "ask_if": ask_if, |
| 198 | "alone": bool(q.get("alone", False)), |
| 199 | } |
| 200 | ) |
| 201 | by_id = {q["id"]: q for q in qs} |
| 202 | for q in qs: |
| 203 | for dep in q["depends_on"]: |
| 204 | if dep not in by_id or dep == q["id"]: |
| 205 | raise SchemaError( |
| 206 | f"question {q['id']!r}: depends on unknown question {dep!r}" |
| 207 | ) |
| 208 | for dep, vals in q["ask_if"].items(): |
| 209 | names = [c[0] for c in by_id[dep]["choices"]] |
| 210 | if any(v not in names for v in vals): |
| 211 | raise SchemaError( |
| 212 | f"question {q['id']!r}: ask_if values for {dep!r} " |
| 213 | f"must be among {names}" |
| 214 | ) |
| 215 | schedule(qs) # refuses a cycle |
| 216 | samples = value.get("samples", "auto") |
| 217 | if samples == "auto": |
| 218 | policy = { |
| 219 | "mode": "auto", |
| 220 | "max": max(1, min(int(value.get("auto_max", 4)), MAX_SAMPLES)), |
| 221 | "threshold": float(value.get("auto_threshold", 0.1)), |
| 222 | } |
| 223 | elif isinstance(samples, int) and samples >= 1: |
| 224 | policy = {"mode": "fixed", "n": min(samples, MAX_SAMPLES)} |
| 225 | else: |
| 226 | raise SchemaError('schema: samples must be a positive count or "auto"') |
| 227 | ask = value.get("ask") |
| 228 | if ask is not None: |
| 229 | if not isinstance(ask, list) or not ask or any(a not in seen for a in ask): |
| 230 | raise SchemaError("schema: ask must list question ids from this schema") |
| 231 | for q in qs: |
| 232 | if q["id"] in ask and any(d not in ask for d in q["depends_on"]): |
| 233 | raise SchemaError( |
| 234 | f"schema: ask names {q['id']!r} but not everything it depends on" |
| 235 | ) |
| 236 | chunk_rows = value.get("chunk_rows") |
| 237 | if chunk_rows is not None and (not isinstance(chunk_rows, int) or chunk_rows < 8): |
| 238 | raise SchemaError("schema: chunk_rows must be an integer of at least 8") |
| 239 | chunk_prompt = value.get("chunk_prompt", "own") |
| 240 | if chunk_prompt not in ("shared", "own"): |
| 241 | raise SchemaError('schema: chunk_prompt must be "shared" or "own"') |
| 242 | sequential = bool(value.get("sequential", False)) |
| 243 | think = value.get("think", 0) |
| 244 | if isinstance(think, bool) or not isinstance(think, int) or not 0 <= think <= 4096: |
| 245 | raise SchemaError("schema: think must be a thought budget in tokens, 0 to 4096") |
| 246 | return { |
| 247 | "questions": qs, |
| 248 | "instructions": value.get("instructions"), |
| 249 | "policy": policy, |
| 250 | "steps": max(1, min(int(value.get("steps", 1)), 8)), |
| 251 | "think": think, |
| 252 | "ask": ask, |
| 253 | "chunk_rows": chunk_rows, |
| 254 | "chunk_prompt": chunk_prompt, |
| 255 | "sequential": sequential, |
| 256 | "format": "lines" if len(qs) <= 10 else "indexed", |
| 257 | } |
| 258 | |
| 259 | |
| 260 | # Answer template shape: (join between questions, what precedes the label, |
| 261 | # reply instruction). A small schema gets "lines", which is readable. |
| 262 | # "indexed" ("0yes 1no") costs three tokens a question against four or five |
| 263 | # and agreed with "lines" on every set tried: 42 booleans, ten 26-way |
| 264 | # choices, twenty 5-level scores. Past ten questions the saved rows keep a |
| 265 | # schema in one read. Two tokens a question, or no id at all, loses |
| 266 | # alignment beyond about twenty questions, because the id ties a label to |
| 267 | # its question. |
| 268 | FORMATS = { |
| 269 | "lines": ( |
| 270 | "\n", |
| 271 | "{id}: ", |
| 272 | 'Reply with one line per question, in this order, formatted as "id: label".', |
| 273 | ), |
| 274 | "indexed": ( |
| 275 | " ", |
| 276 | "{id}", |
| 277 | "Reply on one line with each question's id immediately followed by its " |
| 278 | "label, separated by single spaces.", |
| 279 | ), |
| 280 | } |
| 281 | |
| 282 | |
| 283 | def system_text(schema, chunked=False): |
| 284 | s = ( |
| 285 | "Answer a fixed set of questions about the state the user provides. " |
| 286 | "Each question lists its allowed answers; reply with exactly one label " |
| 287 | "per question.\n" |
| 288 | ) |
| 289 | if schema.get("instructions"): |
| 290 | s += "\n" + str(schema["instructions"]).strip() + "\n" |
| 291 | for q in schema["questions"]: |
| 292 | s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n" |
| 293 | for (name, desc), label in zip(q["choices"], q["labels"]): |
| 294 | if q["type"] == "noul": |
| 295 | s += f" {label}: {str(desc).strip()}\n" if desc else f" {label}\n" |
| 296 | elif desc: |
| 297 | s += f" {label}: {name} ({str(desc).strip()})\n" |
| 298 | else: |
| 299 | s += f" {label}: {name}\n" |
| 300 | s += "\n" + FORMATS[schema.get("format", "lines")][2] |
| 301 | if chunked: |
| 302 | s += ( |
| 303 | " A reply may cover only some of the questions; answer every line " |
| 304 | "that is present." |
| 305 | ) |
| 306 | return s |
| 307 | |
| 308 | |
| 309 | def answer_text(qs, labels, fmt="lines"): |
| 310 | join, lead, _ = FORMATS[fmt] |
| 311 | return join.join( |
| 312 | lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels) |
| 313 | ) |
| 314 | |
| 315 | |
| 316 | def enc(text): |
| 317 | return TOK.encode(text, add_special_tokens=False) |
| 318 | |
| 319 | |
| 320 | def init_tokenizer(tok): |
| 321 | global TOK, SCAFFOLD, THOUGHT_OPEN, THOUGHT_CLOSE |
| 322 | TOK = tok |
| 323 | THOUGHT_OPEN = enc("<|channel>thought\n") |
| 324 | THOUGHT_CLOSE = enc("<channel|>") |
| 325 | SCAFFOLD = enc(SCAFFOLD_TEXT) |
| 326 | assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, ( |
| 327 | "the thought tags must tokenize apart" |
| 328 | ) |
| 329 | |
| 330 | |
| 331 | def resolve_template(qs, head, lead, fmt): |
| 332 | """Tokenize the answer template and find each question's slot. Every label |
| 333 | must change exactly one token, at the same position for all of a question's |
| 334 | labels, or this raises SchemaError. ``head`` is the token run the canvas |
| 335 | starts with: the empty thought block for a plain read, and empty when the |
| 336 | prompt already ends the thought channel. ``lead`` is the text before the |
| 337 | first answer: the join when earlier answers are in the prompt, so the |
| 338 | tokens match one joint template.""" |
| 339 | base_labels = [0] * len(qs) |
| 340 | base = head + enc(lead + answer_text(qs, base_labels, fmt)) |
| 341 | if len(base) + 1 > CANVAS_LEN: |
| 342 | raise SchemaError( |
| 343 | f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}" |
| 344 | ) |
| 345 | if len(qs) == 1 and len(base) + 1 > CANVAS_LEN: |
| 346 | raise SchemaError( |
| 347 | f"question {qs[0]['id']!r} alone needs {len(base) + 1} canvas rows" |
| 348 | ) |
| 349 | slots = [] |
| 350 | for qi, q in enumerate(qs): |
| 351 | pos = None |
| 352 | ids = [0] * len(q["labels"]) |
| 353 | for li in range(1, len(q["labels"])): |
| 354 | labels = list(base_labels) |
| 355 | labels[qi] = li |
| 356 | e = head + enc(lead + answer_text(qs, labels, fmt)) |
| 357 | if len(e) != len(base): |
| 358 | raise SchemaError( |
| 359 | f"question {q['id']!r}: label {q['labels'][li]!r} is not a " |
| 360 | "single token" |
| 361 | ) |
| 362 | diffs = [i for i in range(len(e)) if e[i] != base[i]] |
| 363 | if len(diffs) != 1 or (pos is not None and diffs[0] != pos): |
| 364 | raise SchemaError( |
| 365 | f"question {q['id']!r}: labels do not share one template slot" |
| 366 | ) |
| 367 | pos = diffs[0] |
| 368 | ids[li] = e[pos] |
| 369 | ids[0] = base[pos] |
| 370 | if len(set(ids)) != len(ids): |
| 371 | raise SchemaError( |
| 372 | f"question {q['id']!r}: two labels tokenize to the same id" |
| 373 | ) |
| 374 | slots.append({"pos": pos, "label_ids": ids}) |
| 375 | return base, slots |
| 376 | |
| 377 | |
| 378 | _template_cache = {} |
| 379 | |
| 380 | |
| 381 | def template_for(schema, head, lead): |
| 382 | fmt = schema.get("format", "lines") |
| 383 | key = json.dumps( |
| 384 | [head, lead, fmt] + [(q["id"], q["labels"]) for q in schema["questions"]] |
| 385 | ) |
| 386 | if key not in _template_cache: |
| 387 | _template_cache[key] = resolve_template(schema["questions"], head, lead, fmt) |
| 388 | return _template_cache[key] |
| 389 | |
| 390 | |
| 391 | # ---------------------------------------------------------------------------- |
| 392 | # Reads |
| 393 | # ---------------------------------------------------------------------------- |
| 394 | |
| 395 | |
| 396 | def canvas_width(template): |
| 397 | """Smallest multiple of CANVAS_STEP that holds the template and the turn close.""" |
| 398 | need = len(template) + 1 |
| 399 | return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP) |
| 400 | |
| 401 | |
| 402 | def constrained_xargs(): |
| 403 | """Read over the labels only. The engine runs the unembedding, sampler and |
| 404 | self-conditioning over the request's logprob_token_ids instead of the whole |
| 405 | vocabulary. Same argmax, about a quarter less GPU time per read. Logprobs |
| 406 | come back normalized over the labels.""" |
| 407 | return {"diffusion_constrained": True} if ARGS.constrained else {} |
| 408 | |
| 409 | |
| 410 | def pin_xargs(template, slots, steps): |
| 411 | """Past one denoise step the template must be held, or accept/renoise |
| 412 | rewrites it: pin every canvas position that is not an answer slot.""" |
| 413 | if steps <= 1: |
| 414 | return {} |
| 415 | free = {s["pos"] for s in slots} |
| 416 | return { |
| 417 | "diffusion_pinned": [p for p in range(canvas_width(template)) if p not in free] |
| 418 | } |
| 419 | |
| 420 | |
| 421 | def build_canvas(template, slots, seed): |
| 422 | rng = random.Random(seed) |
| 423 | canvas = list(template) + [TURN_CLOSE] |
| 424 | canvas += [PAD] * (canvas_width(template) - len(canvas)) |
| 425 | for s in slots: |
| 426 | canvas[s["pos"]] = rng.randrange(VOCAB) |
| 427 | return canvas |
| 428 | |
| 429 | |
| 430 | def label_id_union(slots): |
| 431 | ids = sorted({i for s in slots for i in s["label_ids"]}) |
| 432 | return ids[:128] # vLLM's cap per request. A schema needs far fewer. |
| 433 | |
| 434 | |
| 435 | def upstream_chat(body, timeout=600): |
| 436 | req = urllib.request.Request( |
| 437 | ARGS.upstream.rstrip("/") + "/v1/chat/completions", |
| 438 | data=json.dumps(body).encode(), |
| 439 | headers={"content-type": "application/json"}, |
| 440 | ) |
| 441 | return json.load(urllib.request.urlopen(req, timeout=timeout)) |
| 442 | |
| 443 | |
| 444 | def upstream_completions(body, timeout=600): |
| 445 | req = urllib.request.Request( |
| 446 | ARGS.upstream.rstrip("/") + "/v1/completions", |
| 447 | data=json.dumps(body).encode(), |
| 448 | headers={"content-type": "application/json"}, |
| 449 | ) |
| 450 | return json.load(urllib.request.urlopen(req, timeout=timeout)) |
| 451 | |
| 452 | |
| 453 | def chat_prompt_ids(sys_text, state_text, thinking=False): |
| 454 | """The prompt the chat endpoint would build, as token ids, ending after |
| 455 | the model turn marker. Text states only. ``thinking`` turns the chat |
| 456 | template's thinking marker on.""" |
| 457 | messages = [ |
| 458 | {"role": "system", "content": sys_text}, |
| 459 | {"role": "user", "content": state_text}, |
| 460 | ] |
| 461 | out = TOK.apply_chat_template( |
| 462 | messages, tokenize=True, add_generation_prompt=True, enable_thinking=thinking |
| 463 | ) |
| 464 | ids = ( |
| 465 | out["input_ids"] if hasattr(out, "keys") else out |
| 466 | ) # newer transformers return a dict |
| 467 | return [int(t) for t in ids] |
| 468 | |
| 469 | |
| 470 | def think(sys_text, state_text, budget): |
| 471 | """A read prefix that ends with a thought the model wrote: the chat prompt |
| 472 | with thinking on, the open tag, up to ``budget`` generated tokens, the |
| 473 | close tag. Returns the prefix and a diagnostics dict for the thought.""" |
| 474 | prompt = chat_prompt_ids(sys_text, state_text, thinking=True) + THOUGHT_OPEN |
| 475 | started = time.time() |
| 476 | d = upstream_completions( |
| 477 | { |
| 478 | "model": ARGS.model, |
| 479 | "prompt": prompt, |
| 480 | "max_tokens": budget, |
| 481 | "logprobs": 0, |
| 482 | "return_tokens_as_token_ids": True, |
| 483 | "stop_token_ids": THOUGHT_CLOSE, |
| 484 | } |
| 485 | ) |
| 486 | ids = [int(t.split(":")[1]) for t in d["choices"][0]["logprobs"]["tokens"]] |
| 487 | closed = THOUGHT_CLOSE[0] in ids |
| 488 | if closed: |
| 489 | ids = ids[: ids.index(THOUGHT_CLOSE[0])] |
| 490 | info = { |
| 491 | "tokens": len(ids), |
| 492 | "closed": closed, |
| 493 | "ms": (time.time() - started) * 1e3, |
| 494 | "text": TOK.decode(ids), |
| 495 | } |
| 496 | return prompt + ids + THOUGHT_CLOSE, info |
| 497 | |
| 498 | |
| 499 | def think_chat(sys_text, state_content, budget): |
| 500 | """A thought written with the image in view: the chat endpoint with |
| 501 | thinking on, capped at ``budget`` tokens and cut at the close tag. |
| 502 | Returns the thought's token ids and a diagnostics dict.""" |
| 503 | body = { |
| 504 | "model": ARGS.model, |
| 505 | "messages": [ |
| 506 | {"role": "system", "content": sys_text}, |
| 507 | {"role": "user", "content": state_content}, |
| 508 | ], |
| 509 | "max_tokens": budget, |
| 510 | "logprobs": True, |
| 511 | "top_logprobs": 0, |
| 512 | "return_tokens_as_token_ids": True, |
| 513 | "stop_token_ids": THOUGHT_CLOSE, |
| 514 | "chat_template_kwargs": {"enable_thinking": True}, |
| 515 | } |
| 516 | started = time.time() |
| 517 | d = upstream_chat(body) |
| 518 | choice = d["choices"][0] |
| 519 | ids = [ |
| 520 | int(t["token"].split(":")[1]) |
| 521 | for t in (choice.get("logprobs") or {}).get("content") or [] |
| 522 | ] |
| 523 | if ids[: len(THOUGHT_OPEN)] == THOUGHT_OPEN: # the model opens the channel itself |
| 524 | ids = ids[len(THOUGHT_OPEN) :] |
| 525 | closed = THOUGHT_CLOSE[0] in ids |
| 526 | if closed: |
| 527 | ids = ids[: ids.index(THOUGHT_CLOSE[0])] |
| 528 | return ids, { |
| 529 | "tokens": len(ids), |
| 530 | "closed": closed, |
| 531 | "ms": (time.time() - started) * 1e3, |
| 532 | "text": TOK.decode(ids), |
| 533 | } |
| 534 | |
| 535 | |
| 536 | def one_read( |
| 537 | schema, template, slots, sys_text, state_content, seed, prefix=None, thinking=False |
| 538 | ): |
| 539 | if prefix is not None: |
| 540 | return one_read_continuation(schema, template, slots, prefix, seed) |
| 541 | messages = [ |
| 542 | {"role": "system", "content": sys_text}, |
| 543 | {"role": "user", "content": state_content}, |
| 544 | ] |
| 545 | body = { |
| 546 | "model": ARGS.model, |
| 547 | "messages": messages, |
| 548 | "max_tokens": len(template) + 1, |
| 549 | "logprobs": True, |
| 550 | "top_logprobs": TOPK, |
| 551 | # Exact logprobs for every label at every position. With a long |
| 552 | # option list most labels never rank in the top-k, and the model's |
| 553 | # mass sits on tokens that spell the option name instead. |
| 554 | "logprob_token_ids": label_id_union(slots), |
| 555 | "return_tokens_as_token_ids": True, |
| 556 | "chat_template_kwargs": {"enable_thinking": thinking}, |
| 557 | "vllm_xargs": { |
| 558 | "diffusion_seed_canvas": build_canvas(template, slots, seed), |
| 559 | "diffusion_canvas_length": canvas_width(template), |
| 560 | "diffusion_max_steps": schema["steps"], |
| 561 | "diffusion_read_only": True, |
| 562 | **pin_xargs(template, slots, schema["steps"]), |
| 563 | **constrained_xargs(), |
| 564 | }, |
| 565 | } |
| 566 | d = upstream_chat(body) |
| 567 | content = d["choices"][0]["logprobs"]["content"] |
| 568 | out = [] |
| 569 | for q, s in zip(schema["questions"], slots): |
| 570 | top = { |
| 571 | int(t["token"].split(":")[1]): t["logprob"] |
| 572 | for t in content[s["pos"]]["top_logprobs"] |
| 573 | } |
| 574 | out.append(slot_distribution(top, s["label_ids"])) |
| 575 | return out, d.get("usage", {}) |
| 576 | |
| 577 | |
| 578 | def slot_distribution(top, label_ids): |
| 579 | """Label probabilities at one slot from the returned logprobs: every |
| 580 | label's own value plus the argmax token. Read-only logprobs are at |
| 581 | temperature 1, so the label softmax uses them directly. The entropy is |
| 582 | over that returned set.""" |
| 583 | floor = min(top.values()) - 5.0 |
| 584 | lp_t = [top.get(i, floor) for i in label_ids] |
| 585 | mx = max(lp_t) |
| 586 | ex = [math.exp(x - mx) for x in lp_t] |
| 587 | probs = [e / sum(ex) for e in ex] |
| 588 | top_p = [math.exp(v) for v in top.values()] |
| 589 | return { |
| 590 | "probs": probs, |
| 591 | "label_mass": sum(math.exp(x) for x in lp_t), |
| 592 | "entropy": -sum(p * math.log(p) for p in top_p if p > 0), |
| 593 | "argmax_is_label": max(top, key=top.get) in label_ids, |
| 594 | } |
| 595 | |
| 596 | |
| 597 | def one_read_continuation(schema, template, slots, prompt_ids, seed): |
| 598 | """A read whose prompt already holds the thought scaffold and earlier |
| 599 | answer lines, sent as token ids so the chat template cannot alter it.""" |
| 600 | body = { |
| 601 | "model": ARGS.model, |
| 602 | "prompt": prompt_ids, |
| 603 | "max_tokens": len(template) + 1, |
| 604 | "logprobs": TOPK, |
| 605 | "logprob_token_ids": label_id_union(slots), |
| 606 | "return_tokens_as_token_ids": True, |
| 607 | "vllm_xargs": { |
| 608 | "diffusion_seed_canvas": build_canvas(template, slots, seed), |
| 609 | "diffusion_canvas_length": canvas_width(template), |
| 610 | "diffusion_max_steps": schema["steps"], |
| 611 | "diffusion_read_only": True, |
| 612 | **pin_xargs(template, slots, schema["steps"]), |
| 613 | **constrained_xargs(), |
| 614 | }, |
| 615 | } |
| 616 | d = upstream_completions(body) |
| 617 | rows = d["choices"][0]["logprobs"]["top_logprobs"] |
| 618 | out = [] |
| 619 | for q, sl in zip(schema["questions"], slots): |
| 620 | top = {int(k.split(":")[1]): v for k, v in rows[sl["pos"]].items()} |
| 621 | out.append(slot_distribution(top, sl["label_ids"])) |
| 622 | return out, d.get("usage", {}) |
| 623 | |
| 624 | |
| 625 | def read_many( |
| 626 | schema, |
| 627 | template, |
| 628 | slots, |
| 629 | sys_text, |
| 630 | state_content, |
| 631 | seed, |
| 632 | n, |
| 633 | prefix=None, |
| 634 | thinking=False, |
| 635 | ): |
| 636 | results = [None] * n |
| 637 | errors = [None] * n |
| 638 | usages = [None] * n |
| 639 | |
| 640 | def run(k): |
| 641 | try: |
| 642 | results[k], usages[k] = one_read( |
| 643 | schema, |
| 644 | template, |
| 645 | slots, |
| 646 | sys_text, |
| 647 | state_content, |
| 648 | seed + k * 7919, |
| 649 | prefix, |
| 650 | thinking, |
| 651 | ) |
| 652 | except Exception as e: # raised again below |
| 653 | errors[k] = e |
| 654 | |
| 655 | threads = [threading.Thread(target=run, args=(k,)) for k in range(n)] |
| 656 | for t in threads: |
| 657 | t.start() |
| 658 | for t in threads: |
| 659 | t.join() |
| 660 | for e in errors: |
| 661 | if e is not None: |
| 662 | raise e |
| 663 | return results, usages |
| 664 | |
| 665 | |
| 666 | def schedule(qs): |
| 667 | """Questions in stages: a question's stage comes after the stages of |
| 668 | everything it depends on. Declaration order is kept within a stage.""" |
| 669 | ids = {q["id"] for q in qs} |
| 670 | pending = list(qs) |
| 671 | done: set = set() |
| 672 | levels = [] |
| 673 | while pending: |
| 674 | level = [ |
| 675 | q |
| 676 | for q in pending |
| 677 | if all(d in done or d not in ids for d in q["depends_on"]) |
| 678 | ] |
| 679 | if not level: |
| 680 | raise SchemaError( |
| 681 | "schema: dependency cycle among " + ", ".join(q["id"] for q in pending) |
| 682 | ) |
| 683 | levels.append(level) |
| 684 | done |= {q["id"] for q in level} |
| 685 | pending = [q for q in pending if q["id"] not in done] |
| 686 | return levels |
| 687 | |
| 688 | |
| 689 | def chunk_groups(schema, qs): |
| 690 | """``qs`` split, in order, into the fewest groups whose answer templates |
| 691 | fit ``chunk_rows`` (the canvas by default). A question marked alone gets |
| 692 | its own group.""" |
| 693 | limit = schema.get("chunk_rows") or CANVAS_LEN |
| 694 | groups, group = [], [] |
| 695 | for q in qs: |
| 696 | if q["alone"]: |
| 697 | if group: |
| 698 | groups.append(group) |
| 699 | group = [] |
| 700 | groups.append([q]) |
| 701 | continue |
| 702 | trial = group + [q] |
| 703 | rows = ( |
| 704 | len(SCAFFOLD) |
| 705 | + len( |
| 706 | enc(answer_text(trial, [0] * len(trial), schema.get("format", "lines"))) |
| 707 | ) |
| 708 | + 1 |
| 709 | ) |
| 710 | if rows > limit and group: |
| 711 | groups.append(group) |
| 712 | group = [q] |
| 713 | else: |
| 714 | group = trial |
| 715 | if group: |
| 716 | groups.append(group) |
| 717 | return groups |
| 718 | |
| 719 | |
| 720 | def answer_name(q, a): |
| 721 | """The answer as the name ask_if compares against: yes or no, an option |
| 722 | name, or a level name.""" |
| 723 | if a is None: |
| 724 | return None |
| 725 | return a["label"] if q["type"] == "noul" else a.get("choice", a.get("level")) |
| 726 | |
| 727 | |
| 728 | def decide(schema, state_content, seed): |
| 729 | """One decision. Questions run in stages by their dependencies. A stage |
| 730 | is one joint read, chunked by the canvas, with a question marked alone |
| 731 | in its own read. Later stages condition on every earlier answer: a |
| 732 | prefilled continuation for a text state, or the answers restated in the |
| 733 | state for an image. A question whose ask_if condition failed is skipped |
| 734 | and its answer is null.""" |
| 735 | started = time.time() |
| 736 | qs = [ |
| 737 | q |
| 738 | for q in schema["questions"] |
| 739 | if not schema.get("ask") or q["id"] in schema["ask"] |
| 740 | ] |
| 741 | levels = schedule(qs) |
| 742 | text_state = isinstance(state_content, str) |
| 743 | fmt = schema["format"] |
| 744 | join = FORMATS[fmt][0] |
| 745 | # More than one read in sequence needs the full question list in every |
| 746 | # prompt, so that later reads continue one answer. |
| 747 | chained = len(levels) > 1 or schema["sequential"] |
| 748 | sys_full = system_text(schema, chunked=not text_state and chained) |
| 749 | base_ids, thought = None, None |
| 750 | if chained and text_state: |
| 751 | if schema["think"]: |
| 752 | base_ids, thought = think(sys_full, state_content, schema["think"]) |
| 753 | else: |
| 754 | base_ids = chat_prompt_ids(sys_full, state_content) + SCAFFOLD |
| 755 | shared = schema["chunk_prompt"] == "shared" |
| 756 | answered, lines, earlier = {}, [], [] |
| 757 | parts, stages, chunks, skipped = [], [], [], {} |
| 758 | by_id = {q["id"]: q for q in schema["questions"]} |
| 759 | |
| 760 | def run(group, k, conditioned): |
| 761 | sub = dict(schema, questions=group) |
| 762 | # The thought is written once: in the prefix of a chained text decision, |
| 763 | # or in the first read of anything else. |
| 764 | sub["think"] = 0 if (chained and text_state) or conditioned else schema["think"] |
| 765 | if text_state: |
| 766 | if conditioned: |
| 767 | prefix, lead, sys_text = ( |
| 768 | base_ids + enc(join.join(lines)), |
| 769 | join, |
| 770 | sys_full, |
| 771 | ) |
| 772 | elif chained: |
| 773 | prefix, lead, sys_text = (base_ids if thought else None), "", sys_full |
| 774 | else: |
| 775 | prefix, lead, sys_text = ( |
| 776 | None, |
| 777 | "", |
| 778 | (system_text(schema, chunked=True) if shared else system_text(sub)), |
| 779 | ) |
| 780 | return decide_group( |
| 781 | sub, sys_text, state_content, seed + 104729 * k, prefix, lead |
| 782 | ) |
| 783 | state = state_content |
| 784 | if conditioned: |
| 785 | text = next( |
| 786 | (p["text"] for p in state_content if p.get("type") == "text"), "" |
| 787 | ) |
| 788 | text += "\n\nAnswers so far:\n" + "\n".join(earlier) |
| 789 | state = [p for p in state_content if p.get("type") != "text"] + [ |
| 790 | {"type": "text", "text": text} |
| 791 | ] |
| 792 | sys_text = sys_full |
| 793 | else: |
| 794 | sys_text = sys_full if (chained or shared) else system_text(sub) |
| 795 | return decide_group(sub, sys_text, state, seed + 104729 * k) |
| 796 | |
| 797 | def absorb(group, body, rows): |
| 798 | parts.append((body, rows)) |
| 799 | chunks.append([q["id"] for q in group]) |
| 800 | answered.update(body["answers"]) |
| 801 | lines.append( |
| 802 | answer_text( |
| 803 | group, |
| 804 | [q["labels"].index(body["answers"][q["id"]]["label"]) for q in group], |
| 805 | fmt, |
| 806 | ) |
| 807 | ) |
| 808 | earlier.extend( |
| 809 | f"{q['id']}: {answer_name(q, body['answers'][q['id']])}" for q in group |
| 810 | ) |
| 811 | |
| 812 | k = 0 |
| 813 | for level in levels: |
| 814 | asked = [] |
| 815 | for q in level: |
| 816 | failed = next( |
| 817 | ( |
| 818 | (dep, vals) |
| 819 | for dep, vals in q["ask_if"].items() |
| 820 | if answer_name(by_id[dep], answered.get(dep)) not in vals |
| 821 | ), |
| 822 | None, |
| 823 | ) |
| 824 | if failed: |
| 825 | answered[q["id"]] = None |
| 826 | skipped[q["id"]] = { |
| 827 | "because": failed[0], |
| 828 | "was": answer_name(by_id[failed[0]], answered.get(failed[0])), |
| 829 | "wanted": failed[1], |
| 830 | } |
| 831 | continue |
| 832 | asked.append(q) |
| 833 | if not asked: |
| 834 | continue |
| 835 | stages.append([q["id"] for q in asked]) |
| 836 | groups = chunk_groups(schema, asked) |
| 837 | conditioned = bool(lines) |
| 838 | if schema["sequential"] or len(groups) == 1: |
| 839 | for group in groups: |
| 840 | body, rows = run( |
| 841 | group, k, conditioned or (schema["sequential"] and bool(lines)) |
| 842 | ) |
| 843 | absorb(group, body, rows) |
| 844 | k += 1 |
| 845 | else: |
| 846 | with ThreadPoolExecutor(max_workers=min(len(groups), MAX_PARALLEL)) as ex: |
| 847 | results = list( |
| 848 | ex.map( |
| 849 | lambda gk, conditioned=conditioned: run( |
| 850 | gk[1], gk[0], conditioned |
| 851 | ), |
| 852 | [(k + i, g) for i, g in enumerate(groups)], |
| 853 | ) |
| 854 | ) |
| 855 | for group, (body, rows) in zip(groups, results): |
| 856 | absorb(group, body, rows) |
| 857 | k += len(groups) |
| 858 | |
| 859 | answers = {q["id"]: answered.get(q["id"]) for q in qs} |
| 860 | diag_q = {} |
| 861 | for body, _ in parts: |
| 862 | diag_q.update(body["diagnostics"]["questions"]) |
| 863 | # A thought written here (a chained text decision) is outside every |
| 864 | # group's row count. One written inside a group is already counted there. |
| 865 | extra_rows = thought["tokens"] if thought else 0 |
| 866 | if thought is None: |
| 867 | thoughts = [b["diagnostics"].get("thought") for b, _ in parts] |
| 868 | thought = ( |
| 869 | thoughts[0] if len(parts) == 1 else ([t for t in thoughts if t] or None) |
| 870 | ) |
| 871 | one = len(parts) == 1 and not skipped |
| 872 | diagnostics = { |
| 873 | "steps": schema["steps"], |
| 874 | "stages": stages, |
| 875 | "skipped": skipped, |
| 876 | "chunks": chunks, |
| 877 | "chunk_prompt": "full" if chained else schema["chunk_prompt"], |
| 878 | "sequential": schema["sequential"], |
| 879 | "conditioning": ( |
| 880 | None |
| 881 | if len(stages) <= 1 and not schema["sequential"] |
| 882 | else ("prefill" if text_state else "restated") |
| 883 | ), |
| 884 | "thought": thought, |
| 885 | "samples": ( |
| 886 | parts[0][0]["diagnostics"]["samples"] |
| 887 | if one |
| 888 | else { |
| 889 | "n": [b["diagnostics"]["samples"]["n"] for b, _ in parts], |
| 890 | "tops": [b["diagnostics"]["samples"]["tops"] for b, _ in parts], |
| 891 | "policy": [b["diagnostics"]["samples"]["policy"] for b, _ in parts], |
| 892 | } |
| 893 | ), |
| 894 | "timing": { |
| 895 | "total_ms": (time.time() - started) * 1e3, |
| 896 | "reads": sum(b["diagnostics"]["timing"]["reads"] for b, _ in parts), |
| 897 | }, |
| 898 | "prompt_tokens": max( |
| 899 | (b["diagnostics"].get("prompt_tokens") or 0) for b, _ in parts |
| 900 | ) |
| 901 | or None, |
| 902 | "questions": diag_q, |
| 903 | "engine": "vllm", |
| 904 | } |
| 905 | return {"answers": answers, "diagnostics": diagnostics}, sum( |
| 906 | rows for _, rows in parts |
| 907 | ) + extra_rows |
| 908 | |
| 909 | |
| 910 | def decide_group(schema, sys_text, state_content, seed, prefix=None, lead=""): |
| 911 | started = time.time() |
| 912 | thought = None |
| 913 | head = SCAFFOLD if prefix is None else [] |
| 914 | thinking = False |
| 915 | if prefix is None and schema["think"]: |
| 916 | if isinstance(state_content, str): |
| 917 | prefix, thought = think(sys_text, state_content, schema["think"]) |
| 918 | head = [] |
| 919 | else: |
| 920 | # With images the model writes the thought with the image in view, and |
| 921 | # the server seeds it into the canvas ahead of the answer, so the read |
| 922 | # keeps the image. The canvas bounds the thought. |
| 923 | answer = enc( |
| 924 | answer_text( |
| 925 | schema["questions"], |
| 926 | [0] * len(schema["questions"]), |
| 927 | schema.get("format", "lines"), |
| 928 | ) |
| 929 | ) |
| 930 | fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer) |
| 931 | if fits < 8: |
| 932 | raise SchemaError( |
| 933 | f"think: the canvas leaves {fits} rows for a thought " |
| 934 | "beside this template" |
| 935 | ) |
| 936 | ids, thought = think_chat( |
| 937 | sys_text, state_content, min(schema["think"], fits) |
| 938 | ) |
| 939 | thought["budget"] = min(schema["think"], fits) |
| 940 | head = THOUGHT_OPEN + ids + THOUGHT_CLOSE |
| 941 | thinking = True |
| 942 | template, slots = template_for(schema, head, lead) |
| 943 | policy = schema["policy"] |
| 944 | if policy["mode"] == "fixed": |
| 945 | reads, usages = read_many( |
| 946 | schema, |
| 947 | template, |
| 948 | slots, |
| 949 | sys_text, |
| 950 | state_content, |
| 951 | seed, |
| 952 | policy["n"], |
| 953 | prefix, |
| 954 | thinking, |
| 955 | ) |
| 956 | extended = None |
| 957 | first_entropy = None |
| 958 | else: |
| 959 | reads, usages = read_many( |
| 960 | schema, template, slots, sys_text, state_content, seed, 1, prefix, thinking |
| 961 | ) |
| 962 | first_entropy = { |
| 963 | q["id"]: r["entropy"] for q, r in zip(schema["questions"], reads[0]) |
| 964 | } |
| 965 | extended = ( |
| 966 | max(first_entropy.values()) > policy["threshold"] and policy["max"] > 1 |
| 967 | ) |
| 968 | if extended: |
| 969 | more, more_usages = read_many( |
| 970 | schema, |
| 971 | template, |
| 972 | slots, |
| 973 | sys_text, |
| 974 | state_content, |
| 975 | seed + 1, |
| 976 | policy["max"] - 1, |
| 977 | prefix, |
| 978 | thinking, |
| 979 | ) |
| 980 | reads += more |
| 981 | usages += more_usages |
| 982 | prompt_tokens = next( |
| 983 | (u["prompt_tokens"] for u in usages if u and u.get("prompt_tokens")), None |
| 984 | ) |
| 985 | elapsed_ms = (time.time() - started) * 1e3 |
| 986 | |
| 987 | answers = {} |
| 988 | diag_q = {} |
| 989 | n = len(reads) |
| 990 | for qi, q in enumerate(schema["questions"]): |
| 991 | per = [r[qi]["probs"] for r in reads] |
| 992 | mean = [sum(p[i] for p in per) / n for i in range(len(q["labels"]))] |
| 993 | top = max(range(len(mean)), key=lambda i: mean[i]) |
| 994 | a = { |
| 995 | "type": q["type"], |
| 996 | "label": q["labels"][top], |
| 997 | "confidence": mean[top], |
| 998 | "probabilities": {c[0]: m for c, m in zip(q["choices"], mean)}, |
| 999 | } |
| 1000 | if q["type"] == "noul": |
| 1001 | a["noul"] = mean[0] |
| 1002 | elif q["type"] == "choice": |
| 1003 | a["choice"] = q["choices"][top][0] |
| 1004 | else: |
| 1005 | a["score"] = sum((i + 1) * m for i, m in enumerate(mean)) |
| 1006 | a["level"] = q["choices"][top][0] |
| 1007 | if n > 1: |
| 1008 | var = sum((p[top] - mean[top]) ** 2 for p in per) / (n - 1) |
| 1009 | a["stderr"] = (var / n) ** 0.5 |
| 1010 | a["agreement"] = ( |
| 1011 | sum(1 for p in per if max(range(len(p)), key=lambda i: p[i]) == top) / n |
| 1012 | ) |
| 1013 | answers[q["id"]] = a |
| 1014 | diag_q[q["id"]] = { |
| 1015 | "pos": slots[qi]["pos"], |
| 1016 | "entropy": [r[qi]["entropy"] for r in reads], |
| 1017 | "label_mass": reads[0][qi]["label_mass"], |
| 1018 | "argmax_is_label": reads[0][qi]["argmax_is_label"], |
| 1019 | } |
| 1020 | tops = [ |
| 1021 | { |
| 1022 | q["id"]: [ |
| 1023 | q["labels"][ |
| 1024 | max(range(len(r[qi]["probs"])), key=lambda i: r[qi]["probs"][i]) |
| 1025 | ], |
| 1026 | max(r[qi]["probs"]), |
| 1027 | r[qi]["entropy"], |
| 1028 | ] |
| 1029 | for qi, q in enumerate(schema["questions"]) |
| 1030 | } |
| 1031 | for r in reads |
| 1032 | ] |
| 1033 | return { |
| 1034 | "answers": answers, |
| 1035 | "diagnostics": { |
| 1036 | "steps": schema["steps"], |
| 1037 | "samples": { |
| 1038 | "n": n, |
| 1039 | "tops": tops, |
| 1040 | "policy": dict( |
| 1041 | policy, extended=extended, first_read_entropy=first_entropy |
| 1042 | ), |
| 1043 | }, |
| 1044 | "timing": {"total_ms": elapsed_ms, "reads": n}, |
| 1045 | "thought": thought, |
| 1046 | "prompt_tokens": prompt_tokens, |
| 1047 | "questions": diag_q, |
| 1048 | "engine": "vllm", |
| 1049 | }, |
| 1050 | }, len(template) + 1 + (thought["tokens"] if thought else 0) |
| 1051 | |
| 1052 | |
| 1053 | # ---------------------------------------------------------------------------- |
| 1054 | # Jev's contract |
| 1055 | # ---------------------------------------------------------------------------- |
| 1056 | |
| 1057 | JEV_EXTENSIONS = ( |
| 1058 | "instructions", |
| 1059 | "samples", |
| 1060 | "auto_max", |
| 1061 | "auto_threshold", |
| 1062 | "steps", |
| 1063 | "think", |
| 1064 | "ask", |
| 1065 | "chunk_rows", |
| 1066 | "chunk_prompt", |
| 1067 | "sequential", |
| 1068 | ) |
| 1069 | |
| 1070 | |
| 1071 | def jev_schema(body): |
| 1072 | """This server's schema from a Jev request body.""" |
| 1073 | qs = body.get("questions") |
| 1074 | if not isinstance(qs, dict) or not qs: |
| 1075 | raise SchemaError("questions: needs a non-empty map of id -> question") |
| 1076 | out = [] |
| 1077 | for qid, q in qs.items(): |
| 1078 | if not isinstance(q, dict): |
| 1079 | raise SchemaError(f"question {qid!r}: must be an object") |
| 1080 | kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "") |
| 1081 | item = { |
| 1082 | "id": qid, |
| 1083 | "type": kind, |
| 1084 | "instructions": ins if isinstance(ins, str) else json.dumps(ins), |
| 1085 | } |
| 1086 | if kind == "noul": |
| 1087 | if crit is not None and not isinstance(crit, dict): |
| 1088 | raise SchemaError( |
| 1089 | f"question {qid!r}: noul criteria must be an object with " |
| 1090 | "true and false" |
| 1091 | ) |
| 1092 | item["criteria"] = crit |
| 1093 | elif kind == "choice": |
| 1094 | if not isinstance(crit, dict) or not crit: |
| 1095 | raise SchemaError( |
| 1096 | f"question {qid!r}: choice criteria must map option names " |
| 1097 | "to descriptions" |
| 1098 | ) |
| 1099 | item["options"] = [ |
| 1100 | {"name": str(n), "description": d} for n, d in crit.items() |
| 1101 | ] |
| 1102 | elif kind == "score": |
| 1103 | if not isinstance(crit, list): |
| 1104 | raise SchemaError( |
| 1105 | f"question {qid!r}: score criteria must be an ordered list " |
| 1106 | "of levels" |
| 1107 | ) |
| 1108 | item["levels"] = crit |
| 1109 | else: |
| 1110 | raise SchemaError(f"question {qid!r}: unknown type {kind!r}") |
| 1111 | for key in ("depends_on", "ask_if", "alone"): |
| 1112 | if key in q: |
| 1113 | item[key] = q[key] |
| 1114 | out.append(item) |
| 1115 | schema = {k: body[k] for k in JEV_EXTENSIONS if k in body} |
| 1116 | schema["questions"] = out |
| 1117 | return parse_schema(schema) |
| 1118 | |
| 1119 | |
| 1120 | def image_part(content_type, data): |
| 1121 | return { |
| 1122 | "type": "image_url", |
| 1123 | "image_url": { |
| 1124 | "url": f"data:{content_type};base64," + base64.b64encode(data).decode() |
| 1125 | }, |
| 1126 | } |
| 1127 | |
| 1128 | |
| 1129 | def jev_images(value): |
| 1130 | """Image parts from the body's "images": data URLs, or objects with |
| 1131 | content_type and base64.""" |
| 1132 | parts = [] |
| 1133 | for i, im in enumerate(value or []): |
| 1134 | if isinstance(im, str) and im.startswith("data:image/"): |
| 1135 | parts.append({"type": "image_url", "image_url": {"url": im}}) |
| 1136 | elif ( |
| 1137 | isinstance(im, dict) |
| 1138 | and str(im.get("content_type", "")).startswith("image/") |
| 1139 | and isinstance(im.get("base64"), str) |
| 1140 | ): |
| 1141 | parts.append( |
| 1142 | { |
| 1143 | "type": "image_url", |
| 1144 | "image_url": { |
| 1145 | "url": f"data:{im['content_type']};base64,{im['base64']}" |
| 1146 | }, |
| 1147 | } |
| 1148 | ) |
| 1149 | else: |
| 1150 | raise SchemaError( |
| 1151 | f"images[{i}]: a data:image/... URL or an object with " |
| 1152 | "content_type and base64" |
| 1153 | ) |
| 1154 | return parts |
| 1155 | |
| 1156 | |
| 1157 | def jev_state(body, image_parts=()): |
| 1158 | """The user message: the state as text (as given, or as JSON), with any |
| 1159 | images ahead of it.""" |
| 1160 | state = body.get("state") |
| 1161 | if state is None: |
| 1162 | raise SchemaError("state: required") |
| 1163 | text = state if isinstance(state, str) else json.dumps(state) |
| 1164 | if not image_parts: |
| 1165 | return text |
| 1166 | return list(image_parts) + [{"type": "text", "text": text}] |
| 1167 | |
| 1168 | |
| 1169 | def jev_answer(q, a): |
| 1170 | if a is None: |
| 1171 | return None |
| 1172 | if q["type"] == "noul": |
| 1173 | return {"type": "noul", "noul": a["noul"]} |
| 1174 | if q["type"] == "choice": |
| 1175 | return { |
| 1176 | "type": "choice", |
| 1177 | "choice": a["choice"], |
| 1178 | "probabilities": a["probabilities"], |
| 1179 | "confidence": a["confidence"], |
| 1180 | } |
| 1181 | names = [c[0] for c in q["choices"]] |
| 1182 | probs = {str(i): a["probabilities"][n] for i, n in enumerate(names)} |
| 1183 | return { |
| 1184 | "type": "score", |
| 1185 | "score": sum(i * p for i, p in enumerate(probs.values())), |
| 1186 | "legend": {str(i): n for i, n in enumerate(names)}, |
| 1187 | "probabilities": probs, |
| 1188 | "confidence": a["confidence"], |
| 1189 | } |
| 1190 | |
| 1191 | |
| 1192 | # ---------------------------------------------------------------------------- |
| 1193 | # HTTP |
| 1194 | # ---------------------------------------------------------------------------- |
| 1195 | |
| 1196 | |
| 1197 | def message_text(m): |
| 1198 | c = m.get("content", "") |
| 1199 | if isinstance(c, list): |
| 1200 | return "".join(p.get("text", "") for p in c if isinstance(p, dict)) |
| 1201 | return c if isinstance(c, str) else "" |
| 1202 | |
| 1203 | |
| 1204 | class Handler(BaseHTTPRequestHandler): |
| 1205 | def log_message(self, fmt, *args): |
| 1206 | pass |
| 1207 | |
| 1208 | def _json(self, code, obj): |
| 1209 | body = json.dumps(obj).encode() |
| 1210 | self.send_response(code) |
| 1211 | self.send_header("content-type", "application/json") |
| 1212 | self.send_header("content-length", str(len(body))) |
| 1213 | self.end_headers() |
| 1214 | self.wfile.write(body) |
| 1215 | |
| 1216 | def do_GET(self): |
| 1217 | if self.path == "/health": |
| 1218 | return self._json(200, {"status": "ok"}) |
| 1219 | return self._json(404, {"error": {"message": "unknown route"}}) |
| 1220 | |
| 1221 | def _read_request(self): |
| 1222 | """-> (body, image parts): a JSON body, or multipart/form-data with the |
| 1223 | JSON in a part named request and each image as a file part, in order.""" |
| 1224 | raw = self.rfile.read(int(self.headers.get("content-length", "0"))) |
| 1225 | ctype = self.headers.get("content-type", "") |
| 1226 | if not ctype.lower().startswith("multipart/form-data"): |
| 1227 | return json.loads(raw), [] |
| 1228 | from email.parser import BytesParser |
| 1229 | from email.policy import HTTP |
| 1230 | |
| 1231 | msg = BytesParser(policy=HTTP).parsebytes( |
| 1232 | b"Content-Type: " + ctype.encode() + b"\r\n\r\n" + raw |
| 1233 | ) |
| 1234 | body, images = None, [] |
| 1235 | for part in msg.iter_parts(): |
| 1236 | name = part.get_param("name", header="content-disposition") |
| 1237 | data = part.get_payload(decode=True) |
| 1238 | if name == "request": |
| 1239 | body = json.loads(data) |
| 1240 | elif part.get_content_type().startswith("image/"): |
| 1241 | images.append(image_part(part.get_content_type(), data)) |
| 1242 | else: |
| 1243 | raise ValueError( |
| 1244 | f"part {name!r}: neither the request JSON nor an image" |
| 1245 | ) |
| 1246 | if body is None: |
| 1247 | raise ValueError( |
| 1248 | "multipart needs a part named request holding the JSON body" |
| 1249 | ) |
| 1250 | return body, images |
| 1251 | |
| 1252 | def do_POST(self): |
| 1253 | if API_KEY and self.headers.get("authorization", "") != f"Bearer {API_KEY}": |
| 1254 | return self._json( |
| 1255 | 401, |
| 1256 | { |
| 1257 | "error": { |
| 1258 | "message": "missing or wrong bearer token", |
| 1259 | "type": "authentication_error", |
| 1260 | } |
| 1261 | }, |
| 1262 | ) |
| 1263 | if self.path == "/v1/raw/chat/completions": |
| 1264 | return self._raw_chat() |
| 1265 | try: |
| 1266 | req, images = self._read_request() |
| 1267 | except Exception as e: |
| 1268 | return self._json( |
| 1269 | 400, |
| 1270 | { |
| 1271 | "error": { |
| 1272 | "message": f"invalid body: {e}", |
| 1273 | "type": "invalid_request_error", |
| 1274 | } |
| 1275 | }, |
| 1276 | ) |
| 1277 | if self.path == "/v1/systemone": |
| 1278 | return self._systemone(req, images) |
| 1279 | if self.path == "/v1/chat/completions": |
| 1280 | return self._chat(req) |
| 1281 | return self._json(404, {"error": {"message": "unknown route"}}) |
| 1282 | |
| 1283 | def _raw_chat(self): |
| 1284 | """Pass the body and status through to vLLM's chat completions.""" |
| 1285 | raw = self.rfile.read(int(self.headers.get("content-length", "0"))) |
| 1286 | req = urllib.request.Request( |
| 1287 | ARGS.upstream.rstrip("/") + "/v1/chat/completions", |
| 1288 | data=raw, |
| 1289 | headers={ |
| 1290 | "content-type": self.headers.get("content-type", "application/json") |
| 1291 | }, |
| 1292 | ) |
| 1293 | try: |
| 1294 | with urllib.request.urlopen(req, timeout=600) as r: |
| 1295 | code, body = r.status, r.read() |
| 1296 | except urllib.error.HTTPError as e: |
| 1297 | code, body = e.code, e.read() |
| 1298 | self.send_response(code) |
| 1299 | self.send_header("content-type", "application/json") |
| 1300 | self.send_header("content-length", str(len(body))) |
| 1301 | self.end_headers() |
| 1302 | self.wfile.write(body) |
| 1303 | |
| 1304 | def _decide(self, schema, state, seed): |
| 1305 | """-> (status, body) with the error body already shaped.""" |
| 1306 | try: |
| 1307 | return 200, decide(schema, state, seed) |
| 1308 | except SchemaError as e: |
| 1309 | return 422, {"error": {"message": str(e), "type": "validation_error"}} |
| 1310 | except urllib.error.HTTPError as e: |
| 1311 | return 502, { |
| 1312 | "error": { |
| 1313 | "message": ( |
| 1314 | f"upstream {e.code}: {e.read()[:300].decode(errors='replace')}" |
| 1315 | ), |
| 1316 | "type": "server_error", |
| 1317 | } |
| 1318 | } |
| 1319 | except Exception as e: |
| 1320 | return 500, {"error": {"message": repr(e), "type": "server_error"}} |
| 1321 | |
| 1322 | def _systemone(self, req, images): |
| 1323 | try: |
| 1324 | schema = jev_schema(req) |
| 1325 | state = jev_state(req, images + jev_images(req.get("images"))) |
| 1326 | except SchemaError as e: |
| 1327 | return self._json( |
| 1328 | 422, {"error": {"message": str(e), "type": "validation_error"}} |
| 1329 | ) |
| 1330 | code, result = self._decide(schema, state, int(req.get("seed", 42))) |
| 1331 | if code != 200: |
| 1332 | return self._json(code, result) |
| 1333 | body, completion_tokens = result |
| 1334 | answers = { |
| 1335 | q["id"]: jev_answer(q, body["answers"][q["id"]]) |
| 1336 | for q in schema["questions"] |
| 1337 | } |
| 1338 | labels = " ".join( |
| 1339 | f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items() |
| 1340 | ) |
| 1341 | print( |
| 1342 | f"systemone: {labels} " |
| 1343 | f"reads={body['diagnostics']['timing']['reads']} " |
| 1344 | f"{body['diagnostics']['timing']['total_ms']:.0f}ms", |
| 1345 | flush=True, |
| 1346 | ) |
| 1347 | self._json( |
| 1348 | 200, |
| 1349 | { |
| 1350 | "model": ARGS.model, |
| 1351 | "answers": answers, |
| 1352 | # vLLM's prompt count when the reads reported one, since it covers |
| 1353 | # images. Otherwise the tokenizer's count of the text prompt. |
| 1354 | "usage": { |
| 1355 | "input_tokens": body["diagnostics"].get("prompt_tokens") |
| 1356 | or ( |
| 1357 | len(chat_prompt_ids(system_text(schema), state)) |
| 1358 | if isinstance(state, str) |
| 1359 | else 0 |
| 1360 | ), |
| 1361 | "output_tokens": completion_tokens, |
| 1362 | }, |
| 1363 | "diagnostics": body["diagnostics"], |
| 1364 | }, |
| 1365 | ) |
| 1366 | |
| 1367 | def _chat(self, req): |
| 1368 | msgs = req.get("messages") or [] |
| 1369 | if ( |
| 1370 | len(msgs) != 2 |
| 1371 | or msgs[0].get("role") not in ("system", "developer") |
| 1372 | or msgs[1].get("role") != "user" |
| 1373 | ): |
| 1374 | return self._json( |
| 1375 | 400, |
| 1376 | { |
| 1377 | "error": { |
| 1378 | "message": ( |
| 1379 | "a structured request is exactly two messages: the " |
| 1380 | "schema (system) and the state JSON (user)" |
| 1381 | ), |
| 1382 | "type": "invalid_request_error", |
| 1383 | } |
| 1384 | }, |
| 1385 | ) |
| 1386 | try: |
| 1387 | schema_value = json.loads(message_text(msgs[0])) |
| 1388 | schema = parse_schema(schema_value) |
| 1389 | content = msgs[1].get("content", "") |
| 1390 | has_image = isinstance(content, list) and any( |
| 1391 | isinstance(p, dict) and p.get("type") in ("image_url", "image") |
| 1392 | for p in content |
| 1393 | ) |
| 1394 | if has_image: |
| 1395 | # image parts pass through to vLLM unchanged, with text parts as context |
| 1396 | state = content |
| 1397 | else: |
| 1398 | state = message_text(msgs[1]).strip() |
| 1399 | json.loads(state) |
| 1400 | except SchemaError as e: |
| 1401 | return self._json( |
| 1402 | 400, {"error": {"message": str(e), "type": "invalid_request_error"}} |
| 1403 | ) |
| 1404 | except Exception as e: |
| 1405 | return self._json( |
| 1406 | 400, |
| 1407 | { |
| 1408 | "error": { |
| 1409 | "message": ( |
| 1410 | "system must be a JSON question schema and user " |
| 1411 | f"must be JSON state or image parts: {e}" |
| 1412 | ), |
| 1413 | "type": "invalid_request_error", |
| 1414 | } |
| 1415 | }, |
| 1416 | ) |
| 1417 | code, result = self._decide(schema, state, int(req.get("seed", 42))) |
| 1418 | if code != 200: |
| 1419 | if code == 422: |
| 1420 | result["error"]["type"] = "invalid_request_error" |
| 1421 | code = 400 |
| 1422 | return self._json(code, result) |
| 1423 | body, completion_tokens = result |
| 1424 | content = json.dumps(body, indent=2) |
| 1425 | labels = " ".join( |
| 1426 | f"{k}={v['label'] if v else 'skipped'}" for k, v in body["answers"].items() |
| 1427 | ) |
| 1428 | print( |
| 1429 | f"structured: {labels} " |
| 1430 | f"reads={body['diagnostics']['timing']['reads']} " |
| 1431 | f"{body['diagnostics']['timing']['total_ms']:.0f}ms", |
| 1432 | flush=True, |
| 1433 | ) |
| 1434 | self._json( |
| 1435 | 200, |
| 1436 | { |
| 1437 | "id": f"chatcmpl-{int(time.time() * 1000)}", |
| 1438 | "object": "chat.completion", |
| 1439 | "created": int(time.time()), |
| 1440 | "model": req.get("model", "dgemma-structured"), |
| 1441 | "choices": [ |
| 1442 | { |
| 1443 | "index": 0, |
| 1444 | "message": {"role": "assistant", "content": content}, |
| 1445 | "finish_reason": "stop", |
| 1446 | } |
| 1447 | ], |
| 1448 | "usage": { |
| 1449 | "prompt_tokens": 0, |
| 1450 | "completion_tokens": completion_tokens, |
| 1451 | "total_tokens": completion_tokens, |
| 1452 | }, |
| 1453 | }, |
| 1454 | ) |
| 1455 | |
| 1456 | |
| 1457 | def self_signed(cert_dir): |
| 1458 | """Paths of a self-signed certificate and key in cert_dir, made with |
| 1459 | openssl on first use.""" |
| 1460 | os.makedirs(cert_dir, exist_ok=True) |
| 1461 | cert, key = os.path.join(cert_dir, "djev.crt"), os.path.join(cert_dir, "djev.key") |
| 1462 | if not (os.path.exists(cert) and os.path.exists(key)): |
| 1463 | subprocess.run( |
| 1464 | [ |
| 1465 | "openssl", |
| 1466 | "req", |
| 1467 | "-x509", |
| 1468 | "-newkey", |
| 1469 | "rsa:2048", |
| 1470 | "-nodes", |
| 1471 | "-days", |
| 1472 | "3650", |
| 1473 | "-subj", |
| 1474 | "/CN=djev", |
| 1475 | "-keyout", |
| 1476 | key, |
| 1477 | "-out", |
| 1478 | cert, |
| 1479 | ], |
| 1480 | check=True, |
| 1481 | capture_output=True, |
| 1482 | ) |
| 1483 | return cert, key |
| 1484 | |
| 1485 | |
| 1486 | def serve_tls(host, port, cert_dir): |
| 1487 | cert, key = self_signed(cert_dir) |
| 1488 | ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) |
| 1489 | ctx.load_cert_chain(cert, key) |
| 1490 | srv = ThreadingHTTPServer((host, port), Handler) |
| 1491 | srv.socket = ctx.wrap_socket(srv.socket, server_side=True) |
| 1492 | threading.Thread(target=srv.serve_forever, daemon=True).start() |
| 1493 | return srv |
| 1494 | |
| 1495 | |
| 1496 | def main(): |
| 1497 | global ARGS, CANVAS_LEN, CANVAS_STEP |
| 1498 | p = argparse.ArgumentParser() |
| 1499 | p.add_argument("--upstream", default="http://127.0.0.1:8010") |
| 1500 | p.add_argument("--model", default="dgemma") |
| 1501 | p.add_argument("--tokenizer", default="/models/dgemma", help="HF id or local path") |
| 1502 | p.add_argument("--canvas", type=int, default=64, help="the served canvas length") |
| 1503 | p.add_argument( |
| 1504 | "--canvas-step", |
| 1505 | type=int, |
| 1506 | default=16, |
| 1507 | help="request widths round up to a multiple of this", |
| 1508 | ) |
| 1509 | p.add_argument( |
| 1510 | "--no-constrained", |
| 1511 | dest="constrained", |
| 1512 | action="store_false", |
| 1513 | help="read over the whole vocabulary instead of the labels", |
| 1514 | ) |
| 1515 | p.add_argument("--host", default="0.0.0.0") |
| 1516 | p.add_argument("--port", type=int, default=8011) |
| 1517 | p.add_argument( |
| 1518 | "--tls-port", type=int, default=0, help="also listen with HTTPS here (0 = off)" |
| 1519 | ) |
| 1520 | p.add_argument( |
| 1521 | "--cert-dir", |
| 1522 | default=os.path.expanduser("~/.cache/djev"), |
| 1523 | help="directory for the self-signed certificate", |
| 1524 | ) |
| 1525 | ARGS = p.parse_args() |
| 1526 | CANVAS_LEN = ARGS.canvas |
| 1527 | CANVAS_STEP = ARGS.canvas_step |
| 1528 | init_tokenizer(AutoTokenizer.from_pretrained(ARGS.tokenizer)) |
| 1529 | if ARGS.tls_port: |
| 1530 | serve_tls(ARGS.host, ARGS.tls_port, ARGS.cert_dir) |
| 1531 | print( |
| 1532 | f"structured server https on {ARGS.host}:{ARGS.tls_port} (self-signed)", |
| 1533 | flush=True, |
| 1534 | ) |
| 1535 | print( |
| 1536 | f"structured server on {ARGS.host}:{ARGS.port} -> {ARGS.upstream} " |
| 1537 | f"(canvas {CANVAS_LEN})", |
| 1538 | flush=True, |
| 1539 | ) |
| 1540 | ThreadingHTTPServer((ARGS.host, ARGS.port), Handler).serve_forever() |
| 1541 | |
| 1542 | |
| 1543 | if __name__ == "__main__": |
| 1544 | main() |