Spaces:
Running on Zero
Running on Zero
| """Single-forward-pass typed-decision inference for the demo Space. | |
| Takes a `state` (free-form JSON, optionally containing embedded images) and a `questions` dict | |
| (noul / choice / score primitives) and returns calibrated probabilities -- without generating any | |
| text. The model never explains itself; it reads out logits for a small, fixed set of answer | |
| tokens after one forward pass over the rendered prompt. | |
| This is a standalone reimplementation scoped to exactly what this demo needs. It does not import | |
| or vendor any third-party inference harness -- just `transformers`/`torch` plus the small amount | |
| of prompt-rendering and calibration arithmetic below. | |
| """ | |
| from __future__ import annotations | |
| import base64 | |
| import io | |
| import json | |
| import math | |
| import os | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from transformers import AutoModelForCausalLM, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer | |
| REPO = "nutrientdocs/nutrient-document-decision-private" | |
| SUBFOLDER = "merged" | |
| SYSTEM_PROMPT = ( | |
| "You are a document decision model. You read the Evidence and answer each Criterion " | |
| "by choosing exactly one of the listed options. You never explain. You answer with " | |
| "the single option label only." | |
| ) | |
| IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" | |
| LETTERS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" | |
| # --------------------------------------------------------------------------- model + calibration | |
| _model = None | |
| _tokenizer = None | |
| _processor = None | |
| _multimodal = False | |
| _calibration: dict[str, Any] = {"temperature": {"noul": 1.0, "choice": 1.0, "score": 1.0}, "head": None} | |
| def load(device: str = "cuda"): | |
| global _model, _tokenizer, _processor, _multimodal, _calibration | |
| token = os.environ.get("HF_TOKEN") | |
| _tokenizer = AutoTokenizer.from_pretrained(REPO, subfolder=SUBFOLDER, token=token) | |
| try: | |
| _processor = AutoProcessor.from_pretrained(REPO, subfolder=SUBFOLDER, token=token) | |
| _multimodal = True | |
| except Exception as e: | |
| print(f"[serve] AutoProcessor failed to load, falling back to text-only: {e!r}", flush=True) | |
| _processor = None | |
| _multimodal = False | |
| loader = AutoModelForImageTextToText if _multimodal else AutoModelForCausalLM | |
| _model = loader.from_pretrained(REPO, subfolder=SUBFOLDER, dtype=torch.bfloat16, device_map=device, token=token) | |
| _model.eval() | |
| cal_path = hf_hub_download(REPO, "calibration.json", token=token) | |
| _calibration = json.load(open(cal_path)) | |
| # --------------------------------------------------------------------------- image handling | |
| def _is_image_ref(x: Any) -> bool: | |
| return isinstance(x, dict) and x.get("type") == "image" and isinstance(x.get("source"), str) | |
| def _has_images(state: Any) -> bool: | |
| if _is_image_ref(state): | |
| return True | |
| if isinstance(state, dict): | |
| return any(_has_images(v) for v in state.values()) | |
| if isinstance(state, list): | |
| return any(_has_images(v) for v in state) | |
| return False | |
| def _load_image(source: str) -> Image.Image: | |
| if source.startswith("data:"): | |
| header, _, payload = source.partition(",") | |
| blob = base64.b64decode(payload) | |
| elif source.startswith(("http://", "https://")): | |
| import httpx | |
| blob = httpx.get(source, timeout=30, follow_redirects=True).content | |
| else: | |
| blob = open(source, "rb").read() | |
| return Image.open(io.BytesIO(blob)).convert("RGB") | |
| def _split_images(state: Any) -> tuple[Any, list[Image.Image]]: | |
| """Replace every {"type": "image", "source": ...} with a placeholder; collect the images | |
| in document order.""" | |
| images: list[Image.Image] = [] | |
| def walk(x): | |
| if _is_image_ref(x): | |
| images.append(_load_image(x["source"])) | |
| return IMAGE_PLACEHOLDER | |
| if isinstance(x, dict): | |
| return {k: walk(v) for k, v in x.items()} | |
| if isinstance(x, list): | |
| return [walk(v) for v in x] | |
| return x | |
| return walk(state), images | |
| def _render_text(x: Any) -> str: | |
| if x is None: | |
| return "" | |
| if isinstance(x, str): | |
| return x | |
| return json.dumps(x, ensure_ascii=False, indent=2) | |
| # --------------------------------------------------------------------------- prompt construction | |
| def _prefix(state: Any) -> str: | |
| body = f"# Evidence\n{_render_text(state)}\n\n" | |
| return f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\n{body}" | |
| def _branch_tail(body: str) -> str: | |
| return f"{body}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n" | |
| def _options_block(instructions: Any, labelled: list[tuple[str, str]], ask: str) -> str: | |
| lines = [f"# Criterion\n{_render_text(instructions)}\n", "# Options"] | |
| for label, desc in labelled: | |
| lines.append(f"{label}. {desc}" if desc else f"{label}.") | |
| lines.append(f"\n{ask}\n") | |
| return "\n".join(lines) | |
| def _build_branch(qid: str, q: dict) -> tuple[str, list[str], list[Any]]: | |
| """Returns (branch_text, label_tokens_in_prompt_order, keys_in_same_order).""" | |
| kind = q["type"] | |
| if kind == "noul": | |
| criteria = q.get("criteria") or {} | |
| yes_text = _render_text(criteria.get("true")) or "The statement is true." | |
| no_text = _render_text(criteria.get("false")) or "The statement is false." | |
| pair = [("A", "yes: " + yes_text), ("B", "no: " + no_text)] | |
| body = _options_block(q["instructions"], pair, "Respond with only the letter.") | |
| return _branch_tail(body), ["A", "B"], [True, False] | |
| if kind == "choice": | |
| keys = list(q["criteria"].keys()) | |
| labels = [LETTERS[i] for i in range(len(keys))] | |
| labelled = [] | |
| for lab, k in zip(labels, keys): | |
| desc = _render_text(q["criteria"][k]) | |
| labelled.append((lab, f"{k}: {desc}" if desc else k)) | |
| body = _options_block(q["instructions"], labelled, "Respond with only the letter of the best option.") | |
| return _branch_tail(body), labels, keys | |
| if kind == "score": | |
| keys = list(range(len(q["criteria"]))) | |
| labels = [LETTERS[i] for i in range(len(keys))] | |
| labelled = [(lab, f"(level {k} of {len(keys) - 1}) {_render_text(q['criteria'][k])}") for lab, k in zip(labels, keys)] | |
| body = _options_block(q["instructions"], labelled, "Respond with only the letter of the level that best matches.") | |
| return _branch_tail(body), labels, keys | |
| raise ValueError(f"unknown question type {kind!r}") | |
| # --------------------------------------------------------------------------- calibration | |
| def _features(kind: str, logits: np.ndarray, state_tokens: int) -> np.ndarray: | |
| z = logits - logits.max() | |
| p = np.exp(z) / np.exp(z).sum() | |
| n = len(z) | |
| ent = float(-(p * np.log(np.clip(p, 1e-12, 1))).sum() / math.log(max(n, 2))) | |
| top = np.sort(z)[::-1] | |
| margin = float(top[0] - top[1]) if n > 1 else 0.0 | |
| return np.array([ | |
| 1.0, kind == "noul", kind == "choice", kind == "score", | |
| math.log(n), math.log(max(state_tokens, 1)) / 10.0, ent, min(margin, 20.0) / 10.0, | |
| ]) | |
| def _temperature(kind: str, logits: np.ndarray, state_tokens: int) -> float: | |
| head = _calibration.get("head") | |
| if head is not None: | |
| f = _features(kind, logits, state_tokens) | |
| return float(np.exp(np.clip(f @ np.asarray(head), math.log(0.2), math.log(20.0)))) | |
| return float(_calibration.get("temperature", {}).get(kind, 1.0)) | |
| def _softmax(logits: np.ndarray, temperature: float) -> np.ndarray: | |
| z = logits / max(temperature, 1e-6) | |
| z = z - z.max() | |
| e = np.exp(z) | |
| return e / e.sum() | |
| def _confidence(p: np.ndarray) -> float: | |
| n = len(p) | |
| if n <= 1: | |
| return 1.0 | |
| p = np.clip(p, 1e-12, 1.0) | |
| h = -(p * np.log(p)).sum() | |
| return float(max(0.0, min(1.0, 1.0 - h / math.log(n)))) | |
| # --------------------------------------------------------------------------- forward pass | |
| def _label_id(label: str) -> int: | |
| ids = _tokenizer.encode(label, add_special_tokens=False) | |
| if len(ids) != 1: | |
| raise ValueError(f"label {label!r} is not a single token") | |
| return ids[0] | |
| def _forward(prefix_text: str, images: list[Image.Image], branch_text: str) -> np.ndarray: | |
| """One full forward pass over prefix+branch; returns the last-token logits (fp32 numpy).""" | |
| text = prefix_text + branch_text | |
| if images and not _multimodal: | |
| raise RuntimeError( | |
| "This request includes an image, but the multimodal processor failed to load " | |
| "at startup (see the [serve] log line) -- refusing to silently answer from text alone." | |
| ) | |
| if images: | |
| inputs = _processor(text=[text], images=images, return_tensors="pt") | |
| else: | |
| inputs = _tokenizer(text, return_tensors="pt") | |
| inputs = {k: v.to(_model.device) for k, v in inputs.items()} | |
| out = _model(**inputs, use_cache=False) | |
| return out.logits[0, -1].float().cpu().numpy() | |
| def answer(state: Any, questions: dict[str, dict]) -> dict: | |
| """Mirrors the SystemOneResponse shape: {qid: {type, ...}}""" | |
| rendered_state, images = _split_images(state) | |
| prefix_text = _prefix(rendered_state) | |
| state_tokens = len(_tokenizer.encode(prefix_text, add_special_tokens=False)) | |
| results = {} | |
| for qid, q in questions.items(): | |
| branch_text, labels, keys = _build_branch(qid, q) | |
| row = _forward(prefix_text, images, branch_text) | |
| restricted = np.array([row[_label_id(lb)] for lb in labels]) | |
| kind = q["type"] | |
| temp = _temperature(kind, restricted, state_tokens) | |
| probs = _softmax(restricted, temp) | |
| key_probs = dict(zip(keys, probs.tolist())) | |
| if kind == "noul": | |
| results[qid] = {"type": "noul", "noul": round(key_probs[True], 6)} | |
| elif kind == "choice": | |
| p = {str(k): round(v, 6) for k, v in key_probs.items()} | |
| best = max(p, key=p.get) | |
| results[qid] = { | |
| "type": "choice", "choice": best, "probabilities": p, | |
| "confidence": round(_confidence(np.array(list(p.values()))), 6), | |
| } | |
| else: # score | |
| levels = sorted(key_probs.keys()) | |
| pvec = np.array([key_probs[i] for i in levels]) | |
| score = float((pvec * np.array(levels)).sum()) | |
| results[qid] = { | |
| "type": "score", "score": round(score, 6), | |
| "legend": {str(i): _render_text(q["criteria"][i]) for i in levels}, | |
| "probabilities": {str(i): round(float(key_probs[i]), 6) for i in levels}, | |
| "confidence": round(_confidence(pvec), 6), | |
| } | |
| return results | |