"""In-process /v1/systemone scoring with the joint schema head (shared by serve_head.py and eval_head.py). Response JSON follows the jev-adapter contract (protocol.reduce_probabilities / service.score), so the unchanged runners (for example `jev_adapter.benchmarks.run`) work as is: answers[q] = {type, probabilities | noul, choice, confidence, score, legend, logprobs?} usage = {input_tokens, output_tokens: 0} metadata.evaluations = questions x permutations (the runners' per-question accounting; the head needs only `prefills` = permutations forward passes for the whole request), plus prefill_ms / head_ms. Option permutations (options.permutations = K) rotate choice and noul options (display position p shows original (p + r) mod n, as the adapter's rotation_order); score levels always stay in ascending order. """ import json import math import time from pathlib import Path import torch from render import EncodeError, entropy_confidence, request_from_systemone, token_roles CONFIDENCE_METHOD = "1 - normalized_entropy" class RequestError(ValueError): def __init__(self, code, message, field=None, status=422): super().__init__(message) self.code, self.field, self.status = code, field, status def validate_body(body): """Light structural validation (the adapter's pydantic schema is used when importable).""" try: from jev_adapter.protocol import SystemOneRequest # noqa: F401 SystemOneRequest.model_validate(body) except ImportError: pass except Exception as e: # pydantic ValidationError errors = getattr(e, "errors", lambda: [{"msg": str(e), "loc": ()}])() first = errors[0] if errors else {"msg": str(e), "loc": ()} raise RequestError("invalid_request", first.get("msg", str(e)), ".".join(str(p) for p in first.get("loc", ()))) from None if not isinstance(body, dict) or not isinstance(body.get("questions"), dict) or not body["questions"]: raise RequestError("invalid_request", "questions are required", "questions") class HeadScorer: def __init__(self, model, encoder, served_model, pad_id, temperature_by_type=None, max_concurrency=1, max_batch_tokens=32768, preview=True): self.model, self.encoder, self.served_model, self.pad_id = model, encoder, served_model, pad_id self.preview = bool(preview) # must match training (head config "preview"; --no-preview ablation) self.temperature_by_type = dict(temperature_by_type or {}) self.max_concurrency = max_concurrency self.max_batch_tokens = max_batch_tokens self.device = model.device() def plan(self, body): validate_body(body) req = request_from_systemone(body) opts = body.get("options") or {} k = int(opts.get("permutations", 1)) encs = [] for r in range(k): orders = {} for key, q in req["questions"].items(): n = len(q["options"]) if q["type"] != "score" and r: orders[key] = [(p + r) % n for p in range(n)] try: encs.append(self.encoder.encode(req, "W0", list(req["questions"]), orders, preview=self.preview)) except EncodeError as e: raise RequestError(e.reason, str(e), "questions" if e.reason != "too_long" else "state") from None return req, opts, encs @torch.no_grad() def run(self, encs): """Forward the encodings (one padded micro-batch); returns (per-encoding logits lists, prefill_s, head_s).""" from model import make_batch dev = self.device sync = (lambda: torch.cuda.synchronize()) if dev.type == "cuda" else (lambda: None) batch = make_batch(encs, self.pad_id, dev) sync() t0 = time.perf_counter() states = self.model.backbone_states(batch) sync() t1 = time.perf_counter() outs = [] with torch.autocast(device_type=dev.type, enabled=False): for i, enc in enumerate(encs): n = enc["n_tokens"] taps = [s[i, :n].float() for s in states] roles = torch.tensor(token_roles(enc), device=dev) outs.append([z.float().cpu() for z in self.model.head.forward_one(taps, enc, roles)]) sync() t2 = time.perf_counter() return outs, t1 - t0, t2 - t1 def respond(self, req, opts, encs, outs, prefill_s, head_s, started): temperature = float(opts.get("temperature", 1.0)) if opts.get("temperature_scaling", True) else 1.0 return_logprobs = bool(opts.get("return_logprobs", False)) answers = {} keys = list(req["questions"]) for key in keys: q = req["questions"][key] names = [o["name"] for o in q["options"]] t_type = self.temperature_by_type.get(q["type"], 1.0) probs_sum = [0.0] * len(names) raw = [] for enc, out in zip(encs, outs): z = out[enc["keys"].index(key)].double() logp = torch.log_softmax(z, -1) raw.append({n: float(v) for n, v in zip(names, logp.tolist())}) p = torch.softmax(z / (t_type * temperature), -1).tolist() for j, v in enumerate(p): probs_sum[j] += v total = math.fsum(probs_sum) probs = [v / total for v in probs_sum] ans = {"type": q["type"]} if q["type"] == "noul": ans["noul"] = probs[0] else: ans["probabilities"] = dict(zip(names, probs)) ans["confidence"] = entropy_confidence(probs) if q["type"] == "choice": ans["choice"] = names[max(range(len(names)), key=probs.__getitem__)] else: ans["score"] = math.fsum(i * p for i, p in enumerate(probs)) ans["legend"] = {o["name"]: o.get("description") for o in q["options"]} if return_logprobs: ans["logprobs"] = raw answers[key] = ans input_tokens = sum(e["n_tokens"] for e in encs) usage = {"input_tokens": input_tokens, "output_tokens": 0} if req["images"]: usage["input_tokens_details"] = {"image_tokens": sum(e["image_tokens"] for e in encs)} return {"model": self.served_model, "answers": answers, "usage": usage, "metadata": {"confidence_method": CONFIDENCE_METHOD, "temperature": temperature, **({"temperature_by_type": self.temperature_by_type} if self.temperature_by_type else {}), "evaluations": len(keys) * len(encs), "prefills": len(encs), "service_max_concurrency": self.max_concurrency, "label_scheme": "schema-v1", "cached_tokens": 0, "prefill_ms": prefill_s * 1000, "head_ms": head_s * 1000, "adapter_elapsed_ms": (time.perf_counter() - started) * 1000}} def score(self, body): started = time.perf_counter() if body.get("model") not in (None, self.served_model, "jev-latest"): raise RequestError("model_not_found", "Requested model is not configured.", "model", 404) req, opts, encs = self.plan(body) outs, ps, hs = self.run(encs) return self.respond(req, opts, encs, outs, ps, hs, started) def load_scorer(snapshot=None, head_dir=None, adapter=None, tiny=None, tokenizer=None, served_model=None, temperature_file=None, device="auto", merge=True, max_concurrency=1, memory_cap_bytes=0): """Backbone (+ LoRA adapter merged) + trained head -> HeadScorer. memory_cap_bytes > 0 caps the torch allocator (set_per_process_memory_fraction) before the model is loaded.""" from model import build_model, load_head_weights from render import Encoder dev = torch.device(("cuda" if torch.cuda.is_available() else "cpu") if device == "auto" else device) if dev.type == "cuda" and memory_cap_bytes: total = torch.cuda.get_device_properties(0).total_memory torch.cuda.set_per_process_memory_fraction(min(1.0, float(memory_cap_bytes) / total), 0) cfg = json.loads((Path(head_dir) / "schema_head_config.json").read_text()) if head_dir else {} hc = cfg.get("head", {}) taps = cfg.get("taps") or hc.get("taps") or ["final", 25] overrides = {k: hc[k] for k in ("d", "heads", "ffn", "route_blocks", "interact_blocks", "ord_hidden", "max_nodes") if k in hc} overrides["dropout"] = 0.0 overrides["grad_checkpoint"] = False model, info = build_model(snapshot, tiny=tiny, lora_rank=0, taps=taps, device=dev, grad_checkpointing=False, init_adapter=adapter, head_overrides=overrides) if adapter and merge and model.peft_model is not None: model.backbone = model.peft_model.merge_and_unload() model.peft_model = None if head_dir: load_head_weights(model.head, head_dir) model.eval() if dev.type == "cuda": model.backbone.to(torch.bfloat16) encoder = Encoder.from_snapshot(tokenizer or snapshot, with_processor=True) pad = encoder.tokenizer.pad_token_id if encoder.tokenizer.pad_token_id is not None else encoder.tokenizer.eos_token_id temps = json.loads(Path(temperature_file).read_text()).get("temperature_by_type") if temperature_file else None return HeadScorer(model, encoder, served_model or "standardthinking/standard-schema-8b", pad, temps, max_concurrency=max_concurrency, preview=cfg.get("preview", True)), info