add serve_decide.py

serve_decide.py

@@ -0,0 +1,430 @@
+#!/usr/bin/env python3
+"""vLLM OpenAI-compatible server with a built-in System 1 endpoint (POST /v1/decide) and adaptive thinking.
+
+One engine serves both systems, for JEV-27B / JEV-27B-VL (Qwen3.8) and GEV-26B-A4B (Gemma-4):
+  System 2  the unmodified base model, through the usual OpenAI endpoints (served model name)
+  System 1  the LoRA module "jev-decision" (backbone LoRA + decision head as an lm_head LoRA), through /v1/decide
+
+POST /v1/decide
+  kind       "noul" (yes/no), "score" (0-5) or "choice" (2-256 options)
+  state      string, JSON object, or a list mixing text and images: ["text", {"image": "https://... | data:..."}]
+  question   string
+  options    list of strings (choice only)
+  strategy   choices with more than 16 options: "single" (one pass, labels A-P then Q-Z, AA, ...), "tournament"
+             (groups of <=16 + a final of 16), "permute" (single pass over 4 option orders, averaged); default per model
+  thinking   "off" (default: System 1 only), "auto" (think only when the leading option is below `threshold`),
+             "on" (always think)
+  threshold  System 1 confidence below which "auto" switches thinking on (default per model)
+  reasoning controls, as for the base model's own chat API:
+    chat_template_kwargs  passed to the base model's chat template (e.g. Qwen3.8: {"reasoning_effort": "low"})
+    reasoning_effort      shorthand for chat_template_kwargs.reasoning_effort, if the base model has it
+                          (Qwen3.8 / JEV-27B: xhigh (default), medium, low; Gemma-4 / GEV: not available)
+    think_budget          maximum thinking tokens; default: no limit beyond the context window
+  return_reasoning    include System 2's reasoning text
+  debug      include the System 1 and System 2 distributions
+
+Adaptive thinking: System 1 gives p1 in one pass. If thinking is switched on, the base model reasons in its thinking
+mode over the same state, question and options; when the thinking channel closes, the answer-letter distribution p2 is
+read in one step, and the result is p = (1 - w) * p1 + w * p2 (default w = 0.5: the fast decision and the reasoning get
+equal weight, which keeps the calibration of System 1; p2 alone is close to one-hot and over-confident).
+
+Response (same shape as the hosted JEV API, plus "thinking" when requested):
+  {"kind", "effective_kind", "options", "probabilities", "choice_index", "choice", "adaptation", "protocol", "model",
+   "usage", "elapsed_seconds", "num_model_requests", "thinking": {...}}
+
+Run it like `python -m vllm.entrypoints.openai.api_server` (same flags), with --lora-modules jev-decision=<adapter dir>
+and --trust-request-chat-template. decision_head.json is read from the adapter directory, calibration.json from its
+parent (override with JEV_DECIDE_CALIBRATION).
+"""
+import asyncio
+import json
+import math
+import os
+import string
+import time
+import zlib
+from typing import Literal
+
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+
+import vllm.entrypoints.launchers.api_server.entry as entry
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+
+LORA = os.environ.get("JEV_DECIDE_LORA", "jev-decision")
+MAX_OPTIONS = 256
+CHUNK = 128  # vLLM caps logprob_token_ids at 128 per request
+PROTOCOL = "jev27-bare-v1"
+PLACEHOLDER = "XQXCONTENTXQX"
+
+# Per-model settings. threshold / mix: adaptive-thinking defaults. GEV: fitted on 1,754 questions outside the Decision
+# Index (CommonsenseQA, OpenBookQA, AQuA-RAT, MedMCQA, LogiQA, StrategyQA). JEV-27B: same defaults, not yet validated.
+PROFILES = {
+    "qwen": {"prefix": "", "image": "<|vision_start|><|image_pad|><|vision_end|>", "end_think": "</think>",
+             "after_think": "\n\n", "strategy": "single", "threshold": 0.8, "mix": 0.5, "thinking": "off"},
+    "gemma": {"prefix": "<bos>", "image": "<|image|>", "end_think": "<channel|>", "after_think": "",
+              "strategy": "tournament", "threshold": 0.8, "mix": 0.5, "thinking": "off"},
+}
+# Read the full distribution: override generation_config defaults (top_k/top_p) that would truncate processed logprobs.
+READ = dict(max_tokens=1, temperature=1.0, top_p=1.0, top_k=0, min_p=0.0, repetition_penalty=1.0,
+            add_special_tokens=False, return_tokens_as_token_ids=True)
+S = {}
+
+
+class DecideRequest(BaseModel):
+    kind: Literal["noul", "score", "choice"]
+    state: str | dict | list = ""
+    question: str
+    options: list[str] | None = None
+    strategy: Literal["auto", "single", "tournament", "permute"] = "auto"
+    thinking: Literal["default", "off", "auto", "on"] = "default"
+    threshold: float | None = None
+    think_budget: int | None = None
+    reasoning_effort: str | None = None
+    chat_template_kwargs: dict | None = None
+    return_reasoning: bool = False
+    debug: bool = False
+    system2_only: bool = False   # skip System 1: think and return System 2's answer distribution (for callers that already have p1)
+
+
+# ------------------------------------------------------------------ setup
+def setup(args):
+    from transformers import AutoConfig, AutoTokenizer
+
+    path = next((m.path for m in (args.lora_modules or []) if m.name == LORA), None)
+    if path is None:
+        raise SystemExit(f"serve_decide: start vLLM with --lora-modules {LORA}=<adapter dir>")
+    head = json.load(open(os.path.join(path, "decision_head.json")))
+    calib = os.environ.get("JEV_DECIDE_CALIBRATION") or os.path.join(os.path.dirname(os.path.abspath(path)), "calibration.json")
+    temps = json.load(open(calib))["per_kind"] if os.path.exists(calib) else {"noul": 1.0, "score": 1.0, "choice": 1.0}
+    temps.setdefault("score", 1.0)
+    src = args.tokenizer or args.model
+    tok = AutoTokenizer.from_pretrained(src, trust_remote_code=args.trust_remote_code)
+    mtype = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code).model_type
+    prof = dict(PROFILES["gemma" if "gemma" in mtype else "qwen"])
+    for k in ("strategy", "threshold", "mix", "thinking"):
+        if os.environ.get(f"JEV_DECIDE_{k.upper()}"):
+            v = os.environ[f"JEV_DECIDE_{k.upper()}"]
+            prof[k] = v if k in ("strategy", "thinking") else float(v)
+
+    def single_token_labels(context):
+        out = []
+        for lab in list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase for b in string.ascii_uppercase]:
+            t = tok.encode(lab, add_special_tokens=False)
+            if len(t) == 1 and t[0] in tok.encode(context.format(lab), add_special_tokens=False):
+                out.append((lab, t[0]))
+            if len(out) == MAX_OPTIONS:
+                break
+        return out
+
+    dec = single_token_labels("x\n{}) y")                      # System 1 option-line labels
+    lo, hi = head["slots"]["ranges"]["choice"]
+    assert [t for _, t in dec[: hi - lo]] == head["verbalizer_ids"][lo:hi], "first labels must be the trained A-P head"
+    base = tok.encode("Answer: (", add_special_tokens=False)   # System 2 answer labels
+    ans = []
+    for lab in list(string.ascii_uppercase) + [a + b for a in string.ascii_uppercase for b in string.ascii_uppercase]:
+        ids = tok.encode(f"Answer: ({lab})", add_special_tokens=False)
+        if ids[: len(base)] == base and len(ids) == len(base) + 2:
+            ans.append((lab, ids[len(base)]))
+        if len(ans) == MAX_OPTIONS:
+            break
+    chat = tok.apply_chat_template([{"role": "user", "content": PLACEHOLDER}], tokenize=False, add_generation_prompt=True,
+                                   enable_thinking=True)
+    pre, post = chat.split(PLACEHOLDER)
+    raw = ("{%- for m in messages -%}{%- if m['content'] is string -%}{{ m['content'] }}{%- else -%}"
+           "{%- for c in m['content'] -%}{%- if c['type'] == 'text' -%}{{ c['text'] }}{%- else -%}" + prof["image"] +
+           "{%- endif -%}{%- endfor -%}{%- endif -%}{%- endfor -%}")
+    names = args.served_model_name
+    S.update(tok=tok, has_effort="reasoning_effort" in (tok.chat_template or ""), think_cache={},
+             max_logprobs=int(getattr(args, "max_logprobs", 20) or 20),
+             head=head, temps=temps, labels=[l for l, _ in dec], label_ids=[t for _, t in dec], ans_labels=ans, prof=prof,
+             think_pre=pre, think_post=post, raw=raw, end_think_id=tok.convert_tokens_to_ids(prof["end_think"]),
+             model=(names[0] if isinstance(names, list) else names) or args.model, model_type=mtype)
+
+
+# ------------------------------------------------------------------ helpers
+def _parts(state):
+    if isinstance(state, str):
+        return [{"type": "text", "text": state}], False
+    if isinstance(state, dict):
+        return [{"type": "text", "text": json.dumps(state, ensure_ascii=False)}], False
+    out, img = [], False
+    for p in state:
+        if isinstance(p, str):
+            out.append({"type": "text", "text": p})
+        elif isinstance(p, dict) and "image" in p:
+            out.append({"type": "image_url", "image_url": {"url": p["image"]}}); img = True
+        elif isinstance(p, dict) and p.get("type") == "image_url":
+            out.append(p); img = True
+        elif isinstance(p, dict) and p.get("type") == "text":
+            out.append(p)
+        else:
+            out.append({"type": "text", "text": json.dumps(p, ensure_ascii=False)})
+    return out, img
+
+
+def _softmax(z):
+    m = max(z); e = [math.exp(x - m) for x in z]; s = sum(e)
+    return [x / s for x in e]
+
+
+def _err(msg, code=400):
+    return JSONResponse({"error": {"message": msg, "type": "BadRequestError", "code": code}}, status_code=code)
+
+
+class Upstream(Exception):
+    def __init__(self, resp):
+        self.resp = resp
+
+
+def _check(out):
+    if hasattr(out, "error"):
+        raise Upstream(out)
+    return out
+
+
+class Ctx:
+    def __init__(self, raw):
+        self.raw, self.requests, self.prompt_tokens, self.completion_tokens = raw, 0, 0, 0
+
+    def count(self, out):
+        self.requests += 1
+        if getattr(out, "usage", None):
+            self.prompt_tokens += out.usage.prompt_tokens or 0
+            self.completion_tokens += out.usage.completion_tokens or 0
+
+
+def _readout_allowed(ids):
+    """Read-out path: restrict the next token to `ids` and return all of them as top logprobs when the server allows that
+    many (--max-logprobs); otherwise ask for `ids` explicitly with logprob_token_ids (chunks of 128). The first path also
+    works with speculative decoding; logprob_token_ids does not in this vLLM build. Both give the same distribution."""
+    return ids if len(ids) <= S["max_logprobs"] else None
+
+
+async def _logprobs(ctx, content, has_img, ids, model, allowed=None):
+    """One-token read-out with a fallback to the other read-out path if the first one fails."""
+    try:
+        return await _logprobs_once(ctx, content, has_img, ids, model, allowed)
+    except Upstream:
+        raise
+    except Exception:
+        other = None if allowed is not None else (ids if len(ids) <= S["max_logprobs"] else None)
+        if other is None and allowed is None:
+            raise
+        return await _logprobs_once(ctx, content, has_img, ids, model, other)
+
+
+async def _logprobs_once(ctx, content, has_img, ids, model, allowed=None):
+    """One-token read-out: {token_id: logprob} for `ids` after the content (text or chat parts)."""
+    lp = {}
+    step = CHUNK if allowed is None else len(ids)
+    for i in range(0, len(ids), step):
+        chunk = ids[i: i + step]
+        kw = dict(logprob_token_ids=chunk) if allowed is None else dict(allowed_token_ids=allowed)
+        if has_img:
+            r = ChatCompletionRequest(model=model, messages=[{"role": "user", "content": content}], chat_template=S["raw"],
+                                      add_generation_prompt=False, logprobs=True, top_logprobs=len(chunk) if allowed else 1, **kw, **READ)
+            out = _check(await ctx.raw.app.state.openai_serving_chat.create_chat_completion(r, ctx.raw)); ctx.count(out)
+            lp.update({int(t.token.split(":")[1]): t.logprob for t in out.choices[0].logprobs.content[0].top_logprobs})
+        else:
+            text = "".join(c["text"] for c in content)
+            r = CompletionRequest(model=model, prompt=text, logprobs=len(chunk) if allowed else 1, **kw, **READ)
+            out = _check(await ctx.raw.app.state.openai_serving_completion.create_completion(r, ctx.raw)); ctx.count(out)
+            lp.update({int(k.split(":")[1]): v for k, v in out.choices[0].logprobs.top_logprobs[0].items()})
+        if allowed is not None:
+            break
+    return lp
+
+
+# ------------------------------------------------------------------ System 1
+async def s1_pass(ctx, kind, parts, has_img, question, opts):
+    """One System 1 pass over `opts` (<=256). Returns probabilities aligned with opts."""
+    head, temps, prof = S["head"], S["temps"], S["prof"]
+    lo, hi = head["slots"]["ranges"][kind]
+    if kind == "choice":
+        n = len(opts); ids = S["label_ids"][:n]
+        bias = [head["bias"][lo + i] if lo + i < hi else 0.0 for i in range(n)]
+        lines = [f"{S['labels'][i]}) {o}" for i, o in enumerate(opts)]
+    else:
+        ids = head["verbalizer_ids"][lo:hi]; bias = head["bias"][lo:hi]; lines = opts
+    content = ([{"type": "text", "text": f"{prof['prefix']}[kind] {kind}\n[state] "}] + parts +
+               [{"type": "text", "text": f"\n[question] {question}\n[options]\n" + "\n".join(lines) + "\n[decision]:"}])
+    lp = await _logprobs(ctx, content, has_img, ids, LORA, allowed=_readout_allowed(ids))
+    return _softmax([(max(lp.get(t, -1e9), -1e9) + b) / temps[kind] for t, b in zip(ids, bias)])
+
+
+def _groups(n, k=16):
+    g = math.ceil(n / k); base, extra = divmod(n, g); out, i = [], 0
+    for j in range(g):
+        size = base + (1 if j < extra else 0); out.append(list(range(i, i + size))); i += size
+    return out
+
+
+async def s1_dist(ctx, kind, parts, has_img, question, opts, strategy):
+    if kind != "choice" or len(opts) <= 16 or strategy == "single":
+        return await s1_pass(ctx, kind, parts, has_img, question, opts)
+    n = len(opts)
+    if strategy == "permute":
+        import random
+        rng = random.Random(zlib.crc32(question.encode()))
+        orders = [list(range(n))] + [rng.sample(range(n), n) for _ in range(3)]
+        res = await asyncio.gather(*(s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in o]) for o in orders))
+        p = [0.0] * n
+        for o, r in zip(orders, res):
+            for i, v in zip(o, r):
+                p[i] += v / len(orders)
+        return p
+    groups = _groups(n)  # tournament: groups of <=16 in the given order (in parallel), then a final of 16
+    parts_g = await asyncio.gather(*(s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in g]) for g in groups))
+    in_group = {o: p for g, ps in zip(groups, parts_g) for o, p in zip(g, ps)}
+    chosen = [max(g, key=lambda o: (in_group[o], -o)) for g in groups]
+    rest = sorted((o for g in groups for o in g if o not in set(chosen)), key=lambda o: (-in_group[o], o))
+    fin = sorted(chosen + rest[: max(0, 16 - len(chosen))])
+    final = dict(zip(fin, await s1_pass(ctx, kind, parts, has_img, question, [opts[i] for i in fin])))
+    group_of = {o: gi for gi, g in enumerate(groups) for o in g}
+    share = [0.0] * len(groups); cap = [0.0] * len(groups)
+    for f in fin:
+        share[group_of[f]] += final[f]; cap[group_of[f]] += in_group[f]
+    among = sum(a * b for a, b in zip(share, cap))
+    p = [final[o] * among if o in final else share[group_of[o]] * in_group[o] for o in range(n)]
+    s = sum(p)
+    return [x / s for x in p]
+
+
+# ------------------------------------------------------------------ System 2 (thinking)
+def _think_frame(kwargs):
+    """Chat-template prefix / suffix around the user turn, thinking on, with the caller's template kwargs (cached)."""
+    key = json.dumps(kwargs or {}, sort_keys=True)
+    if key not in S["think_cache"]:
+        kw = {k: v for k, v in (kwargs or {}).items() if k != "enable_thinking"}
+        chat = S["tok"].apply_chat_template([{"role": "user", "content": PLACEHOLDER}], tokenize=False, add_generation_prompt=True,
+                                            enable_thinking=True, **kw)
+        S["think_cache"][key] = tuple(chat.split(PLACEHOLDER))
+    return S["think_cache"][key]
+
+
+async def s2_dist(ctx, kind, parts, has_img, question, opts, budget, want_text, tmpl_kwargs=None):
+    """The base model thinks over the same input; the answer-letter distribution is read after the thinking channel."""
+    shown = ["Yes (true)", "No (false)"] if kind == "noul" else opts
+    labs = S["ans_labels"][: len(shown)]
+    body = "\n".join(f"({l}) {o}" for (l, _), o in zip(labs, shown))
+    tail = (f"\n\nQuestion: {question}\n\nOptions:\n{body}\n\n"
+            "Think it through carefully, then give your final answer on the last line in the form: Answer: (X)")
+    pre, post = _think_frame(tmpl_kwargs)
+    user = [{"type": "text", "text": pre}] + parts + [{"type": "text", "text": tail + post}]
+    seed = zlib.crc32((question + body).encode()) & 0x7FFFFFFF
+    t0 = time.time()
+    gen = ChatCompletionRequest(model=S["model"], messages=[{"role": "user", "content": user}], chat_template=S["raw"],
+                                add_generation_prompt=False, add_special_tokens=False, max_tokens=budget, seed=seed,  # None: up to the context window
+                                stop_token_ids=[S["end_think_id"]], skip_special_tokens=False)
+    out = _check(await ctx.raw.app.state.openai_serving_chat.create_chat_completion(gen, ctx.raw)); ctx.count(out)
+    thought = out.choices[0].message.content or ""
+    finished = out.choices[0].finish_reason == "stop"
+    ntok = out.usage.completion_tokens if out.usage else None
+    read = user + [{"type": "text", "text": thought + S["prof"]["end_think"] + S["prof"]["after_think"] + "Answer: ("}]
+    ids = [t for _, t in labs]
+    lp = await _logprobs(ctx, read, True, ids, S["model"], allowed=_readout_allowed(ids))
+    p = _softmax([max(lp.get(t, -1e9), -1e9) for t in ids])
+    if kind == "noul":
+        p = [p[1], p[0]]  # (A) yes / (B) no -> [P(false), P(true)]
+    info = {"think_tokens": ntok, "think_seconds": round(time.time() - t0, 3), "finished_within_budget": finished}
+    if want_text:
+        info["reasoning"] = thought.replace("<|channel>thought\n", "").strip()
+    return p, info
+
+
+def _fold(p1, p2, w):
+    return [(1 - w) * x + w * y for x, y in zip(p1, p2)]
+
+
+# ------------------------------------------------------------------ routes
+router = APIRouter()
+
+
+@router.get("/v1/decide/info")
+async def info():
+    p = S["prof"]
+    return {"protocol": PROTOCOL, "model": S["model"], "model_type": S["model_type"], "max_options": len(S["labels"]),
+            "native_choice_options": len(S["head"]["slots"]["verbalizers"]) - 8, "temperatures": S["temps"],
+            "defaults": {"strategy": p["strategy"], "thinking": p["thinking"], "threshold": p["threshold"], "mix": p["mix"],
+                         "think_budget": None},
+            "reasoning_controls": {"enable_thinking": True, "reasoning_effort": ["xhigh", "medium", "low"] if S["has_effort"] else None,
+                                   "think_budget": "max thinking tokens (default: up to the context window)"}}
+
+
+@router.post("/v1/decide")
+async def decide(req: DecideRequest, raw: Request):
+    t0 = time.time()
+    if req.kind == "choice":
+        opts = req.options or []
+        if not 2 <= len(opts) <= len(S["labels"]):
+            return _err(f"choice needs 2-{len(S['labels'])} options, got {len(opts)}")
+    else:
+        opts = ["false", "true"] if req.kind == "noul" else [str(i) for i in range(6)]
+    prof = S["prof"]
+    thinking = prof["thinking"] if req.thinking == "default" else req.thinking
+    if req.kind == "score":
+        if req.thinking in ("auto", "on"):
+            return _err("thinking is supported for noul and choice")
+        thinking = "off"
+    tmpl = dict(req.chat_template_kwargs or {})
+    if req.reasoning_effort is not None:
+        if not S["has_effort"]:
+            return _err("this base model has no reasoning_effort setting; control the reasoning length with think_budget")
+        tmpl["reasoning_effort"] = req.reasoning_effort
+    strategy = prof["strategy"] if req.strategy == "auto" else req.strategy
+    tau = prof["threshold"] if req.threshold is None else req.threshold
+    parts, has_img = _parts(req.state)
+    ctx = Ctx(raw)
+    try:
+        if req.system2_only:
+            if req.kind == "score":
+                return _err("system2_only is supported for noul and choice")
+            p2, think = await s2_dist(ctx, req.kind, parts, has_img, req.question, opts, req.think_budget, req.return_reasoning, tmpl)
+            k = max(range(len(p2)), key=p2.__getitem__)
+            return {"kind": req.kind, "options": opts, "probabilities": p2, "choice_index": k, "choice": opts[k], "system": 2,
+                    "model": S["model"], "thinking": {"used": True, "budget": req.think_budget, **think},
+                    "usage": {"prompt_tokens": ctx.prompt_tokens, "completion_tokens": ctx.completion_tokens},
+                    "elapsed_seconds": time.time() - t0, "num_model_requests": ctx.requests}
+        p1 = await s1_dist(ctx, req.kind, parts, has_img, req.question, opts, strategy)
+        probs, think = p1, None
+        if thinking == "on" or (thinking == "auto" and max(p1) < tau):
+            p2, think = await s2_dist(ctx, req.kind, parts, has_img, req.question, opts, req.think_budget, req.return_reasoning, tmpl)
+            probs = _fold(p1, p2, prof["mix"])
+            think = {"used": True, **think}
+            if req.debug:
+                think.update(system1=p1, system2=p2)
+        elif thinking != "off":
+            think = {"used": False}
+            if req.debug:
+                think["system1"] = p1
+    except Upstream as e:
+        return JSONResponse(e.resp.model_dump(), status_code=e.resp.error.code)
+    k = max(range(len(probs)), key=probs.__getitem__)
+    native = req.kind != "choice" or len(opts) <= 16
+    resp = {"kind": req.kind, "effective_kind": req.kind, "options": opts, "probabilities": probs, "choice_index": k,
+            "choice": opts[k], "adaptation": "native" if native else f"{strategy}", "protocol": PROTOCOL, "model": S["model"],
+            "usage": {"prompt_tokens": ctx.prompt_tokens, "completion_tokens": ctx.completion_tokens,
+                      "total_tokens": ctx.prompt_tokens + ctx.completion_tokens},
+            "elapsed_seconds": time.time() - t0, "num_model_requests": ctx.requests}
+    if think is not None:
+        resp["thinking"] = {"mode": thinking, "threshold": tau, "budget": req.think_budget, **({"chat_template_kwargs": tmpl} if tmpl else {}), **think}
+    return resp
+
+
+_build_app = entry.build_app
+
+
+def build_app(args, *a, **kw):
+    app = _build_app(args, *a, **kw)
+    setup(args)
+    app.include_router(router)
+    return app
+
+
+entry.build_app = build_app
+
+if __name__ == "__main__":
+    entry.main()