Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch) | |
| import html | |
| import json | |
| import random | |
| import re | |
| import time | |
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from huggingface_hub import snapshot_download | |
| from safetensors import safe_open | |
| from transformers import AutoConfig, AutoProcessor | |
| from transformers.activations import ACT2FN | |
| from transformers.cache_utils import DynamicCache | |
| from transformers.initialization import no_init_weights | |
| from transformers.models.diffusion_gemma import modeling_diffusion_gemma as mdg | |
| MODEL_ID = "Singularity-Compute/super-decider" | |
| # --------------------------------------------------------------------------------------------- | |
| # Server settings, matching runtime/Dockerfile of github.com/Singularity-Compute/super-decider | |
| # --------------------------------------------------------------------------------------------- | |
| CANVAS_LEN = 128 # CANVAS=128 | |
| CANVAS_STEP = 16 | |
| VOCAB = 262144 | |
| TURN_CLOSE = 106 | |
| PAD = 0 | |
| TEMPS = {"score": 0.375} # DJEV_TEMPS=score=0.375 | |
| NOUL_COMMIT = 0.8 # DJEV_NOUL_COMMIT=0.8 | |
| ADAPT_GATES = {"choice": 0.7, "score": 0.7, "noul": 0.85} # DJEV_ADAPTIVE_THINK | |
| ADAPT_BUDGET = 512 | |
| ADAPT_MAX_Q = 2 | |
| THINK_STEPS = 8 # DJEV_THINK_STEPS=8 | |
| BIDIR_MAXPROMPT = 8192 # DJEV_BIDIR_MAXPROMPT | |
| MAX_QUESTIONS = 64 | |
| MAX_PROMPT_TOKENS = 32768 # demo cap | |
| MAX_IMAGES = 8 | |
| SCAFFOLD_TEXT = "<|channel>thought\n<channel|>" | |
| # --------------------------------------------------------------------------------------------- | |
| # Model: DiffusionGemma 26B-A4B with ModelOpt NVFP4 experts (W4A16, dequantized on the fly on GPU) | |
| # --------------------------------------------------------------------------------------------- | |
| _FP4 = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] | |
| _LUT2 = torch.tensor([[_FP4[b & 15], _FP4[b >> 4]] for b in range(256)], dtype=torch.bfloat16) | |
| def _dequant(q, s, g, lut2): | |
| """q [n, R, C/2] uint8 (low nibble first), s [n, R, C/16] fp8 block scales, g [n, R, 1] fp32 global.""" | |
| n, r, c2 = q.shape | |
| v = lut2[q.int()].view(n, r, c2 // 8, 16) | |
| sc = s.float() * g | |
| return (v.float() * sc[..., None]).to(torch.bfloat16).view(n, r, c2 * 2) | |
| class NVFP4Experts(nn.Module): | |
| """Drop-in for DiffusionGemmaTextExperts holding NVFP4-packed weights.""" | |
| def __init__(self, config): | |
| super().__init__() | |
| self.num_experts = config.num_experts | |
| self.act_fn = ACT2FN[config.hidden_activation] | |
| # tiny placeholders so the stock _init_weights isinstance branch is harmless; removed after build | |
| self.gate_up_proj = nn.Parameter(torch.zeros(1)) | |
| self.down_proj = nn.Parameter(torch.zeros(1)) | |
| self.register_buffer("lut2", _LUT2.clone(), persistent=False) | |
| def forward(self, hidden_states, top_k_index, top_k_weights): | |
| out = torch.zeros_like(hidden_states) | |
| with torch.no_grad(): | |
| mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) # [E, K, T] | |
| hit = torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero().flatten() | |
| if hit.numel() == 0: | |
| return out | |
| gu = _dequant(self.gu_q[hit], self.gu_s[hit], self.gu_g[hit], self.lut2) | |
| dn = _dequant(self.dn_q[hit], self.dn_s[hit], self.dn_g[hit], self.lut2) | |
| for j, e in enumerate(hit.tolist()): | |
| k_pos, tok = torch.where(mask[e]) | |
| x = hidden_states[tok] | |
| gate, up = F.linear(x, gu[j]).chunk(2, dim=-1) | |
| h = F.linear(self.act_fn(gate) * up, dn[j]) | |
| h = h * top_k_weights[tok, k_pos, None] | |
| out.index_add_(0, tok, h.to(out.dtype)) | |
| return out | |
| def _set_param(root, path, value): | |
| *mods, leaf = path.split(".") | |
| m = root | |
| for p in mods: | |
| m = getattr(m, p) | |
| setattr(m, leaf, value) | |
| def load_model(): | |
| local = snapshot_download(MODEL_ID) | |
| config = AutoConfig.from_pretrained(local) | |
| if hasattr(config, "quantization_config"): | |
| del config.quantization_config | |
| config.canvas_length = CANVAS_LEN # the served canvas (thought generation) | |
| tc = config.text_config | |
| E, H, I = tc.num_experts, tc.hidden_size, tc.moe_intermediate_size | |
| L = tc.num_hidden_layers | |
| orig = mdg.DiffusionGemmaTextExperts | |
| mdg.DiffusionGemmaTextExperts = NVFP4Experts | |
| try: | |
| with no_init_weights(): | |
| model = mdg.DiffusionGemmaForBlockDiffusion._from_config( | |
| config, dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ) | |
| finally: | |
| mdg.DiffusionGemmaTextExperts = orig | |
| # buffers that are normally computed by _init_weights (rotary tables, embed scales) | |
| for m in model.modules(): | |
| if isinstance(m, (mdg.DiffusionGemmaTextRotaryEmbedding, mdg.DiffusionGemmaTextScaledWordEmbedding)): | |
| model._init_weights(m) | |
| idx = json.load(open(os.path.join(local, "model.safetensors.index.json")))["weight_map"] | |
| files = {f: safe_open(os.path.join(local, f), framework="pt") for f in set(idx.values())} | |
| def get(k): | |
| return files[idx[k]].get_tensor(k) | |
| dec = model.model.decoder | |
| t0 = time.time() | |
| for li in range(L): | |
| ex = dec.layers[li].experts | |
| gu_q = torch.empty(E, 2 * I, H // 2, dtype=torch.uint8) | |
| gu_s = torch.empty(E, 2 * I, H // 16, dtype=torch.float8_e4m3fn) | |
| gu_g = torch.empty(E, 2 * I, 1, dtype=torch.float32) | |
| dn_q = torch.empty(E, H, I // 2, dtype=torch.uint8) | |
| dn_s = torch.empty(E, H, I // 16, dtype=torch.float8_e4m3fn) | |
| dn_g = torch.empty(E, H, 1, dtype=torch.float32) | |
| for e in range(E): | |
| p = f"model.decoder.layers.{li}.experts.{e}." | |
| gu_q[e, :I] = get(p + "gate_proj.weight") | |
| gu_q[e, I:] = get(p + "up_proj.weight") | |
| gu_s[e, :I] = get(p + "gate_proj.weight_scale") | |
| gu_s[e, I:] = get(p + "up_proj.weight_scale") | |
| gu_g[e, :I] = get(p + "gate_proj.weight_scale_2").float() | |
| gu_g[e, I:] = get(p + "up_proj.weight_scale_2").float() | |
| dn_q[e] = get(p + "down_proj.weight") | |
| dn_s[e] = get(p + "down_proj.weight_scale") | |
| dn_g[e] = get(p + "down_proj.weight_scale_2").float() | |
| del ex.gate_up_proj, ex.down_proj | |
| ex.register_buffer("gu_q", gu_q) | |
| ex.register_buffer("gu_s", gu_s) | |
| ex.register_buffer("gu_g", gu_g) | |
| ex.register_buffer("dn_q", dn_q) | |
| ex.register_buffer("dn_s", dn_s) | |
| ex.register_buffer("dn_g", dn_g) | |
| print(f"experts packed in {time.time() - t0:.0f}s", flush=True) | |
| sd = {k: get(k) for k in idx if ".experts." not in k} | |
| res = model.load_state_dict(sd, strict=False) | |
| del sd | |
| # tie the encoder text stack (and lm_head) to the decoder, as in the checkpoint | |
| enc = model.model.encoder.language_model | |
| tied = 0 | |
| for el, dl in zip(enc.layers, dec.layers): | |
| el.experts = dl.experts | |
| for name, p in dl.named_parameters(): | |
| _set_param(el, name, p) | |
| tied += 1 | |
| enc.norm.weight = dec.norm.weight | |
| enc.embed_tokens.weight = dec.embed_tokens.weight | |
| model.lm_head.weight = dec.embed_tokens.weight | |
| missing = [ | |
| k for k in res.missing_keys | |
| if not (k.startswith("model.encoder.language_model.") or k == "lm_head.weight" or ".experts." in k) | |
| ] | |
| print(f"tied {tied} encoder params; unexpected={res.unexpected_keys[:5]} missing={missing[:10]}", flush=True) | |
| files.clear() | |
| return local, model.eval() | |
| LOCAL_DIR, model = load_model() | |
| model.to("cuda") | |
| processor = AutoProcessor.from_pretrained(LOCAL_DIR) | |
| TOK = processor.tokenizer | |
| EOS = int(TOK.eos_token_id) if TOK.eos_token_id is not None else 1 | |
| THOUGHT_OPEN = TOK.encode("<|channel>thought\n", add_special_tokens=False) | |
| THOUGHT_CLOSE = TOK.encode("<channel|>", add_special_tokens=False) | |
| SCAFFOLD = TOK.encode(SCAFFOLD_TEXT, add_special_tokens=False) | |
| assert THOUGHT_OPEN + THOUGHT_CLOSE == SCAFFOLD, "the thought tags must tokenize apart" | |
| SLIDING_WINDOW = model.config.text_config.sliding_window | |
| SOFTCAP = model.config.text_config.final_logit_softcapping | |
| class SchemaError(ValueError): | |
| pass | |
| # --------------------------------------------------------------------------------------------- | |
| # Schema, prompt and answer template (ported from runtime/structured_server.py) | |
| # --------------------------------------------------------------------------------------------- | |
| FORMATS = { | |
| "lines": ("\n", "{id}: ", 'Reply with one line per question, in this order, formatted as "id: label".'), | |
| "indexed": ( | |
| " ", | |
| "{id}", | |
| "Reply on one line with each question's id immediately followed by its label, separated by single spaces.", | |
| ), | |
| } | |
| DECISION_PREAMBLE = ( | |
| "Answer a fixed set of questions about the state the user provides. " | |
| "Each question lists its allowed answers; reply with exactly one label " | |
| "per question.\n" | |
| ) | |
| def enc(text): | |
| return TOK.encode(text, add_special_tokens=False) | |
| def jev_schema(body): | |
| qs = body.get("questions") | |
| if not isinstance(qs, dict) or not qs: | |
| raise SchemaError("questions: needs a non-empty map of id -> question") | |
| if len(qs) > MAX_QUESTIONS: | |
| raise SchemaError(f"at most {MAX_QUESTIONS} questions") | |
| out = [] | |
| for qid, q in qs.items(): | |
| qid = str(qid).strip() | |
| if not qid or ":" in qid or "\n" in qid: | |
| raise SchemaError(f"question id {qid!r} must be non-empty, no ':' or newline") | |
| if not isinstance(q, dict): | |
| raise SchemaError(f"question {qid!r}: must be an object") | |
| if any(k in q for k in ("depends_on", "ask_if")): | |
| raise SchemaError(f"question {qid!r}: depends_on / ask_if are not supported in this demo") | |
| kind, crit, ins = q.get("type"), q.get("criteria"), q.get("instructions", "") | |
| ins = ins if isinstance(ins, str) else json.dumps(ins) | |
| if kind == "noul": | |
| if crit is not None and not isinstance(crit, dict): | |
| raise SchemaError(f"question {qid!r}: noul criteria must be an object with true and false") | |
| crit = crit or {} | |
| choices = [("yes", crit.get("true")), ("no", crit.get("false"))] | |
| labels = ["yes", "no"] | |
| elif kind == "choice": | |
| if not isinstance(crit, dict) or not crit: | |
| raise SchemaError(f"question {qid!r}: choice criteria must map option names to descriptions") | |
| choices = [(str(n), d) for n, d in crit.items()] | |
| labels = [chr(ord("A") + i) for i in range(len(choices))] | |
| elif kind == "score": | |
| if not isinstance(crit, list): | |
| raise SchemaError(f"question {qid!r}: score criteria must be an ordered list of levels") | |
| choices = [(str(level), None) for level in crit] | |
| labels = ( | |
| [str(i + 1) for i in range(len(choices))] | |
| if len(choices) <= 9 | |
| else [chr(ord("A") + i) for i in range(len(choices))] | |
| ) | |
| elif kind in ("span", "spans"): | |
| raise SchemaError(f"question {qid!r}: span questions are not supported in this demo") | |
| else: | |
| raise SchemaError(f"question {qid!r}: unknown type {kind!r} (use noul, choice or score)") | |
| if len(choices) < 2: | |
| raise SchemaError(f"question {qid!r}: needs at least two alternatives") | |
| if len(choices) > 26: | |
| raise SchemaError(f"question {qid!r}: at most 26 alternatives in this demo") | |
| out.append({"id": qid, "type": kind, "instructions": ins, "choices": choices, "labels": labels}) | |
| fmt = "lines" if len(out) <= 10 else "indexed" | |
| schema = {"questions": out, "instructions": body.get("instructions"), "format": fmt} | |
| if fmt == "indexed" and not all(fits_indexed(q) for q in out): | |
| schema["format"] = "lines" | |
| return schema | |
| def fits_indexed(q): | |
| try: | |
| resolve_template([q], [], "", "indexed") | |
| return True | |
| except SchemaError: | |
| return False | |
| def system_text(schema): | |
| s = DECISION_PREAMBLE | |
| if schema.get("instructions"): | |
| s += "\n" + str(schema["instructions"]).strip() + "\n" | |
| for q in schema["questions"]: | |
| s += f"\nQuestion {q['id']}: {q['instructions'].strip()}\n" | |
| for (name, desc), label in zip(q["choices"], q["labels"]): | |
| if q["type"] == "noul": | |
| s += f" {label}: {str(desc).strip()}\n" if desc else f" {label}\n" | |
| elif desc: | |
| s += f" {label}: {name} ({str(desc).strip()})\n" | |
| else: | |
| s += f" {label}: {name}\n" | |
| s += "\n" + FORMATS[schema.get("format", "lines")][2] | |
| return s | |
| def chat_messages(sys_text, state_content): | |
| """DJEV_DOC_FIRST=1: the fixed preamble is the system turn; the questions follow the state.""" | |
| rest = sys_text[len(DECISION_PREAMBLE):].lstrip("\n") | |
| if isinstance(state_content, str): | |
| user = state_content + "\n\n" + rest | |
| else: | |
| user = list(state_content) + [{"type": "text", "text": "\n\n" + rest}] | |
| return [{"role": "system", "content": DECISION_PREAMBLE.rstrip("\n")}, {"role": "user", "content": user}], rest | |
| def answer_text(qs, labels, fmt="lines"): | |
| join, lead, _ = FORMATS[fmt] | |
| return join.join(lead.format(id=q["id"]) + q["labels"][i] for q, i in zip(qs, labels)) | |
| def resolve_template(qs, head, lead, fmt): | |
| base_labels = [0] * len(qs) | |
| base = head + enc(lead + answer_text(qs, base_labels, fmt)) | |
| if len(base) + 1 > CANVAS_LEN: | |
| raise SchemaError(f"answer template is {len(base)} tokens; the canvas holds {CANVAS_LEN - 1}") | |
| slots = [] | |
| for qi, q in enumerate(qs): | |
| pos = None | |
| ids = [0] * len(q["labels"]) | |
| for li in range(1, len(q["labels"])): | |
| labels = list(base_labels) | |
| labels[qi] = li | |
| e = head + enc(lead + answer_text(qs, labels, fmt)) | |
| if len(e) != len(base): | |
| raise SchemaError(f"question {q['id']!r}: label {q['labels'][li]!r} is not a single token") | |
| diffs = [i for i in range(len(e)) if e[i] != base[i]] | |
| if len(diffs) != 1 or (pos is not None and diffs[0] != pos): | |
| raise SchemaError(f"question {q['id']!r}: labels do not share one template slot") | |
| pos = diffs[0] | |
| ids[li] = e[pos] | |
| ids[0] = base[pos] | |
| if len(set(ids)) != len(ids): | |
| raise SchemaError(f"question {q['id']!r}: two labels tokenize to the same id") | |
| slots.append({"pos": pos, "label_ids": ids}) | |
| return base, slots | |
| def chunk_groups(schema): | |
| groups, group = [], [] | |
| for q in schema["questions"]: | |
| trial = group + [q] | |
| rows = len(SCAFFOLD) + len(enc(answer_text(trial, [0] * len(trial), schema["format"]))) + 1 | |
| if rows > CANVAS_LEN and group: | |
| groups.append(group) | |
| group = [q] | |
| else: | |
| group = trial | |
| if group: | |
| groups.append(group) | |
| return groups | |
| def canvas_width(template): | |
| need = len(template) + 1 | |
| return min(CANVAS_LEN, -(-need // CANVAS_STEP) * CANVAS_STEP) | |
| def build_canvas(template, slots, seed): | |
| rng = random.Random(seed) | |
| canvas = list(template) + [TURN_CLOSE] | |
| canvas += [PAD] * (canvas_width(template) - len(canvas)) | |
| for s in slots: | |
| canvas[s["pos"]] = rng.randrange(VOCAB) | |
| return canvas | |
| def bidir_ranges(text, user, rest_len): | |
| """DJEV_BIDIR=qsep: one bidirectional range per question block (question line + its options).""" | |
| e_ = TOK(text, add_special_tokens=False, return_offsets_mapping=True) | |
| ids, offs = e_["input_ids"], e_["offset_mapping"] | |
| if len(ids) > BIDIR_MAXPROMPT: | |
| return ids, None | |
| e = text.rfind(user) + len(user) | |
| a = e - rest_len | |
| rest = user[-rest_len:] | |
| starts = [m.start() for m in re.finditer(r"(?m)^Question \S+: ", rest)] | |
| endq = rest.rfind("\n\nReply with") | |
| endq = endq if endq > 0 else len(rest) | |
| out = [] | |
| for k, s0 in enumerate(starts): | |
| c0, c1 = a + s0, a + (starts[k + 1] - 2 if k + 1 < len(starts) else endq) | |
| jdx = [i for i, (x, y) in enumerate(offs) if x >= c0 and y <= c1 and y > x] | |
| if jdx and jdx[-1] - jdx[0] + 1 == len(jdx) and jdx[-1] > jdx[0]: | |
| out.append((jdx[0], jdx[-1] + 1)) | |
| return ids, (out or None) | |
| def text_prompt(sys_text, state_text, thinking): | |
| msgs, rest = chat_messages(sys_text, state_text) | |
| text = TOK.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=thinking) | |
| ids, ranges = bidir_ranges(text, msgs[1]["content"], len(rest)) | |
| return ids, ranges | |
| def image_prompt(sys_text, state_parts, thinking): | |
| msgs, _ = chat_messages(sys_text, state_parts) | |
| content = [] | |
| for p in msgs[1]["content"]: | |
| if p["type"] == "image": | |
| content.append({"type": "image", "image": p["image"]}) | |
| else: | |
| content.append({"type": "text", "text": p["text"]}) | |
| hf_msgs = [{"role": "system", "content": [{"type": "text", "text": msgs[0]["content"]}]}, | |
| {"role": "user", "content": content}] | |
| inputs = processor.apply_chat_template( | |
| hf_msgs, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", | |
| enable_thinking=thinking, | |
| ) | |
| return inputs | |
| # --------------------------------------------------------------------------------------------- | |
| # GPU reads | |
| # --------------------------------------------------------------------------------------------- | |
| def encoder_masks(n, ranges, device): | |
| i = torch.arange(n, device=device) | |
| causal = i[None, :] <= i[:, None] | |
| window = (i[:, None] - i[None, :]) < SLIDING_WINDOW | |
| full, slide = causal, causal & window | |
| if ranges: | |
| rid = torch.full((n,), -1, device=device, dtype=torch.long) | |
| for k, (a, b) in enumerate(ranges): | |
| rid[a:b] = k | |
| same = (rid[:, None] == rid[None, :]) & (rid[:, None] >= 0) | |
| full, slide = full | same, slide | same | |
| return {"full_attention": full[None, None], "sliding_attention": slide[None, None]} | |
| def encode_prompt(prompt_ids, ranges=None, mm=None, upto=None): | |
| """Prefill the encoder KV cache. Text prompts get the per-question bidirectional mask.""" | |
| upto = len(prompt_ids) if upto is None else upto | |
| ids = torch.tensor([prompt_ids[:upto]], device="cuda") | |
| cache = DynamicCache(config=model.config.get_text_config(decoder=True)) | |
| pos = torch.arange(upto, device="cuda")[None] | |
| kw = {} | |
| if mm is not None: | |
| kw = {k: v.to("cuda") for k, v in mm.items() if k in ("pixel_values", "image_position_ids", "mm_token_type_ids")} | |
| if "pixel_values" in kw: | |
| kw["pixel_values"] = kw["pixel_values"].to(torch.bfloat16) | |
| masks = model.model.encoder.create_masks_for_generate( | |
| config=model.config, | |
| inputs_embeds=torch.empty((1, upto, 0), dtype=torch.bfloat16, device="cuda"), | |
| attention_mask=torch.ones((1, upto), dtype=torch.bool, device="cuda"), | |
| past_key_values=cache, | |
| position_ids=pos, | |
| mm_token_type_ids=kw.get("mm_token_type_ids"), | |
| ) | |
| else: | |
| masks = encoder_masks(upto, [r for r in (ranges or []) if r[1] <= upto], "cuda") | |
| model.model.encoder(input_ids=ids, attention_mask=masks, past_key_values=cache, position_ids=pos, **kw) | |
| return cache | |
| def read(prompt_ids, ranges, template, slots, seed, mm=None): | |
| """One denoise step over a seeded canvas: label probabilities at every answer slot.""" | |
| cache = encode_prompt(prompt_ids, ranges, mm) | |
| n = len(prompt_ids) | |
| canvas = build_canvas(template, slots, seed) | |
| w = len(canvas) | |
| out = model.model.decoder( | |
| decoder_input_ids=torch.tensor([canvas], device="cuda"), | |
| past_key_values=cache, | |
| decoder_position_ids=torch.arange(n, n + w, device="cuda")[None], | |
| ) | |
| hs = out.last_hidden_state[0, [s["pos"] for s in slots]] | |
| logits = model.lm_head(hs).float() | |
| logits = torch.tanh(logits / SOFTCAP) * SOFTCAP | |
| res = [] | |
| for j, s in enumerate(slots): | |
| lp = logits[j, s["label_ids"]] | |
| res.append(torch.softmax(lp, -1).cpu().tolist()) | |
| return res | |
| def write_thought(prompt_ids, ranges, budget, mm=None, seed=42): | |
| """Generate a thought (block diffusion, DJEV_THINK_STEPS denoising steps) until the close tag.""" | |
| torch.manual_seed(seed) | |
| close = THOUGHT_CLOSE[0] | |
| gen_kw = dict( | |
| max_new_tokens=int(budget), max_denoising_steps=THINK_STEPS, eos_token_id=[1, 106, 50, close], | |
| return_dict_in_generate=True, | |
| ) | |
| if mm is None: | |
| cache = encode_prompt(prompt_ids, ranges, None, upto=len(prompt_ids) - 1) | |
| out = model.generate( | |
| input_ids=torch.tensor([prompt_ids[-1:]], device="cuda"), past_key_values=cache, **gen_kw | |
| ) | |
| new = out.sequences[0, 1:].tolist() | |
| else: | |
| kw = {k: v.to("cuda") for k, v in mm.items() if k != "attention_mask"} | |
| if "pixel_values" in kw: | |
| kw["pixel_values"] = kw["pixel_values"].to(torch.bfloat16) | |
| out = model.generate(**kw, **gen_kw) | |
| new = out.sequences[0, mm["input_ids"].shape[1]:].tolist() | |
| if new[: len(THOUGHT_OPEN)] == THOUGHT_OPEN: | |
| new = new[len(THOUGHT_OPEN):] | |
| cut = len(new) | |
| for t in (close, 1, 106, 50): | |
| if t in new: | |
| cut = min(cut, new.index(t)) | |
| return new[:cut] | |
| def summarize(q, probs): | |
| mean = list(probs) | |
| raw_top = max(mean) | |
| t = TEMPS.get(q["type"], 1.0) | |
| if t != 1.0: | |
| w = [max(m, 1e-12) ** (1.0 / t) for m in mean] | |
| mean = [x / sum(w) for x in w] | |
| if NOUL_COMMIT and q["type"] == "noul" and 1 - NOUL_COMMIT < mean[0] < NOUL_COMMIT: | |
| y = NOUL_COMMIT + 0.001 if mean[0] >= 0.5 else 1 - NOUL_COMMIT - 0.001 | |
| mean = [y, 1 - y] | |
| top = max(range(len(mean)), key=lambda i: mean[i]) | |
| names = [c[0] for c in q["choices"]] | |
| if q["type"] == "noul": | |
| ans = {"type": "noul", "noul": mean[0]} | |
| elif q["type"] == "choice": | |
| ans = {"type": "choice", "choice": names[top], "probabilities": dict(zip(names, mean)), | |
| "confidence": mean[top]} | |
| else: | |
| ans = { | |
| "type": "score", | |
| "score": sum(i * p for i, p in enumerate(mean)), | |
| "legend": {str(i): n for i, n in enumerate(names)}, | |
| "probabilities": {str(i): p for i, p in enumerate(mean)}, | |
| "confidence": mean[top], | |
| } | |
| return ans, raw_top | |
| def decide_group(sub, state, seed, think_budget, images_mm_fn): | |
| """One joint read of a question group, optionally preceded by a thought.""" | |
| sys_text = system_text(sub) | |
| thought_ids = None | |
| if isinstance(state, str): | |
| if think_budget: | |
| ids, ranges = text_prompt(sys_text, state, thinking=True) | |
| prompt = ids + THOUGHT_OPEN | |
| thought_ids = write_thought(prompt, ranges, think_budget, seed=seed) | |
| prompt = prompt + thought_ids + THOUGHT_CLOSE | |
| template, slots = resolve_template(sub["questions"], [], "", sub["format"]) | |
| else: | |
| prompt, ranges = text_prompt(sys_text, state, thinking=False) | |
| template, slots = resolve_template(sub["questions"], SCAFFOLD, "", sub["format"]) | |
| if len(prompt) > MAX_PROMPT_TOKENS: | |
| raise SchemaError(f"state too long for this demo ({len(prompt)} tokens > {MAX_PROMPT_TOKENS})") | |
| probs = read(prompt, ranges, template, slots, seed) | |
| n_prompt = len(prompt) | |
| else: | |
| head = SCAFFOLD | |
| thinking = False | |
| if think_budget: | |
| answer = enc(answer_text(sub["questions"], [0] * len(sub["questions"]), sub["format"])) | |
| fits = CANVAS_LEN - 1 - len(THOUGHT_OPEN) - len(THOUGHT_CLOSE) - len(answer) | |
| if fits >= 8: | |
| mm_t = images_mm_fn(sys_text, True) | |
| thought_ids = write_thought(None, None, min(think_budget, fits), mm=mm_t, seed=seed) | |
| head = THOUGHT_OPEN + thought_ids + THOUGHT_CLOSE | |
| thinking = True | |
| mm = images_mm_fn(sys_text, thinking) | |
| template, slots = resolve_template(sub["questions"], head, "", sub["format"]) | |
| prompt = mm["input_ids"][0].tolist() | |
| probs = read(prompt, None, template, slots, seed, mm=mm) | |
| n_prompt = len(prompt) | |
| return probs, n_prompt, (len(template) + 1 + (len(thought_ids) if thought_ids else 0)), thought_ids | |
| def run_decision(body, image_paths, thinking_mode, seed): | |
| schema = jev_schema(body) | |
| state = body.get("state") | |
| if state is None or state == "": | |
| raise SchemaError("state: required") | |
| text = state if isinstance(state, str) else json.dumps(state, ensure_ascii=False) | |
| if image_paths: | |
| from PIL import Image | |
| imgs = [Image.open(p).convert("RGB") for p in image_paths[:MAX_IMAGES]] | |
| state_content = [{"type": "image", "image": im} for im in imgs] + [{"type": "text", "text": text}] | |
| else: | |
| state_content = text | |
| def mm_fn(sys_text, thinking): | |
| return image_prompt(sys_text, state_content, thinking) | |
| always = thinking_mode == "always" | |
| budget = ADAPT_BUDGET if always else 0 | |
| answers, raw, diag = {}, {}, {"reads": 0, "thought_tokens": 0, "adaptive_think": []} | |
| in_tok, out_tok = 0, 0 | |
| groups = chunk_groups(schema) | |
| for k, group in enumerate(groups): | |
| sub = dict(schema, questions=group) | |
| gseed = seed + 104729 * k | |
| probs, n_prompt, n_out, th = decide_group(sub, state_content, gseed, budget, mm_fn) | |
| diag["reads"] += 1 | |
| in_tok = max(in_tok, n_prompt) | |
| out_tok += n_out | |
| if th: | |
| diag["thought_tokens"] += len(th) | |
| for q, p in zip(group, probs): | |
| answers[q["id"]], raw[q["id"]] = summarize(q, p) | |
| # adaptive thinking: re-answer uncertain questions after a thought (text states, <= 2 questions) | |
| if thinking_mode == "auto" and isinstance(state_content, str) and len(group) <= ADAPT_MAX_Q: | |
| low = [q for q in group if raw[q["id"]] < ADAPT_GATES.get(q["type"], 0.0)] | |
| if low: | |
| try: | |
| probs2, _, n2, th2 = decide_group(sub, state_content, gseed, ADAPT_BUDGET, mm_fn) | |
| except SchemaError: | |
| continue | |
| diag["reads"] += 1 | |
| out_tok += n2 | |
| diag["thought_tokens"] += len(th2 or []) | |
| diag["adaptive_think"] += [q["id"] for q in low] | |
| for q, p in zip(group, probs2): | |
| if q in low: | |
| answers[q["id"]], _ = summarize(q, p) | |
| return schema, answers, {"input_tokens": in_tok, "output_tokens": out_tok}, diag | |
| # --------------------------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------------------------- | |
| def render(schema, answers): | |
| cards = [] | |
| for q in schema["questions"]: | |
| a = answers[q["id"]] | |
| names = [c[0] for c in q["choices"]] | |
| if a["type"] == "noul": | |
| rows = [("yes", a["noul"]), ("no", 1 - a["noul"])] | |
| headline = f"P(yes) = {a['noul']:.2f}" | |
| best = "yes" if a["noul"] >= 0.5 else "no" | |
| elif a["type"] == "choice": | |
| rows = [(n, a["probabilities"][n]) for n in names] | |
| headline = f"{html.escape(a['choice'])} · {a['confidence']:.2f}" | |
| best = a["choice"] | |
| else: | |
| rows = [(n, a["probabilities"][str(i)]) for i, n in enumerate(names)] | |
| best = names[max(range(len(names)), key=lambda i: a["probabilities"][str(i)])] | |
| headline = f"score {a['score']:.2f} (0–{len(names) - 1}) · {html.escape(best)}" | |
| bars = "".join( | |
| f"<div class='sd-row'><span class='sd-name{' sd-best' if n == best else ''}'>{html.escape(n)}</span>" | |
| f"<span class='sd-bar'><span style='width:{max(p, 0.003) * 100:.1f}%'></span></span>" | |
| f"<span class='sd-p'>{p:.2f}</span></div>" | |
| for n, p in rows | |
| ) | |
| cards.append( | |
| f"<div class='sd-card'><div class='sd-head'><b>{html.escape(q['id'])}</b>" | |
| f"<span class='sd-type'>{a['type']}</span><span class='sd-hl'>{headline}</span></div>" | |
| f"<div class='sd-q'>{html.escape(q['instructions'])}</div>{bars}</div>" | |
| ) | |
| return "<div class='sd-wrap'>" + "".join(cards) + "</div>" | |
| def _image_paths(images): | |
| paths = [] | |
| for it in images or []: | |
| if isinstance(it, (list, tuple)): | |
| it = it[0] | |
| if isinstance(it, dict): | |
| it = it.get("path") or it.get("name") | |
| if it: | |
| paths.append(it) | |
| return paths | |
| def _gpu_duration(state, questions, thinking="auto", images=None, seed=42, *args, **kwargs): | |
| """Measured on ZeroGPU: one read ~2.5 s, a 512-token thought + read ~24 s.""" | |
| try: | |
| qs = json.loads(questions) if isinstance(questions, str) else questions | |
| if isinstance(qs, dict) and isinstance(qs.get("questions"), dict): | |
| qs = qs["questions"] | |
| n = len(qs) | |
| except Exception: | |
| n = 1 | |
| if thinking == "always" or (thinking == "auto" and n <= ADAPT_MAX_Q and not images): | |
| return 30 | |
| return 10 if n <= 16 else 20 | |
| def decide(state: str, questions: str, thinking: str = "auto", images: list | None = None, seed: int = 42): | |
| """Answer typed decision questions about a state with calibrated probabilities. | |
| Args: | |
| state: The state to decide about: plain text, or JSON (an object or array). | |
| questions: Jev / System One questions as JSON: a map of id -> {"type": "noul"|"choice"|"score", | |
| "instructions": str, "criteria": ...}. choice criteria map option names to descriptions, | |
| score criteria are an ordered list of levels, noul criteria are optional {"true", "false"}. | |
| thinking: "off", "auto" (a short thought only on uncertain questions) or "always". | |
| images: Optional list of image file paths (up to 8) shown ahead of the state. | |
| seed: Seed for the answer-slot noise of the canvas. | |
| Returns: | |
| An HTML summary and the Jev-shaped JSON answer set. | |
| """ | |
| t0 = time.time() | |
| try: | |
| qs = json.loads(questions) if isinstance(questions, str) else questions | |
| if isinstance(qs, dict) and "questions" in qs and isinstance(qs["questions"], dict): | |
| qs = qs["questions"] | |
| except json.JSONDecodeError as e: | |
| raise gr.Error(f"Questions must be valid JSON: {e}") | |
| st = state | |
| if isinstance(state, str) and state.strip()[:1] in ("{", "["): | |
| try: | |
| st = json.loads(state) | |
| except json.JSONDecodeError: | |
| st = state | |
| try: | |
| schema, answers, usage, diag = run_decision( | |
| {"state": st, "questions": qs}, _image_paths(images), thinking or "auto", int(seed) | |
| ) | |
| except SchemaError as e: | |
| raise gr.Error(str(e)) | |
| elapsed = time.time() - t0 | |
| result = { | |
| "model": "super-decider", | |
| "answers": answers, | |
| "usage": usage, | |
| "diagnostics": {**diag, "seconds": round(elapsed, 3)}, | |
| } | |
| print(f"decide: {len(answers)} q, reads={diag['reads']} thought={diag['thought_tokens']} {elapsed:.2f}s", flush=True) | |
| return render(schema, answers), result | |
| EX_TICKET_STATE = ( | |
| 'Order 8812 shipped on Oct 1. The customer writes: "I want to cancel and get my money back." ' | |
| "Policy: shipped orders cannot be cancelled; they can be returned within 30 days of delivery for a full refund." | |
| ) | |
| EX_TICKET_Q = { | |
| "next_step": {"type": "choice", "instructions": "What should support do next?", | |
| "criteria": {"refund_now": "Refund without a return", "start_return": "Start a return", | |
| "escalate": "Hand over to a human"}}, | |
| "allowed": {"type": "noul", "instructions": "Is cancelling this order allowed by the policy?"}, | |
| "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "medium", "high"]}, | |
| } | |
| EX_INVOICE_STATE = "Invoice total 1,240 EUR. The customer paid 2 x 400 EUR and 1 x 340 EUR." | |
| EX_INVOICE_Q = { | |
| "paid": {"type": "noul", "instructions": "Is the invoice fully paid?"}, | |
| "next_step": {"type": "choice", "instructions": "What should accounts receivable do next?", | |
| "criteria": {"close": "Mark as paid and close", "remind": "Send a reminder for the remaining amount", | |
| "escalate": "Hand over to collections"}}, | |
| } | |
| EX_ROUTE_STATE = {"ticket": "I was charged twice for my March invoice.", "plan": "Pro", "customer_since": "2024-02"} | |
| EX_ROUTE_Q = { | |
| "team": {"type": "choice", "instructions": "Which team should handle this ticket?", | |
| "criteria": {"billing": "Payments, invoices, refunds", "technical": "Bugs and outages", | |
| "account": "Login, profile, plan changes", "sales": "New plans and upgrades"}}, | |
| "urgency": {"type": "score", "instructions": "How urgent is this ticket?", "criteria": ["low", "normal", "high"]}, | |
| "churn_risk": {"type": "noul", "instructions": "Does the customer sound likely to cancel?"}, | |
| } | |
| EX_IMAGE_STATE = "Scanned document from the accounts payable inbox." | |
| EX_IMAGE_Q = { | |
| "doc_type": {"type": "choice", "instructions": "What kind of document is this?", | |
| "criteria": {"invoice": "An invoice", "receipt": "A store receipt", "contract": "A contract", | |
| "other": "Something else"}}, | |
| "biggest_line": {"type": "choice", "instructions": "Which line item has the largest amount?", | |
| "criteria": {"paper": "Paper A4 (box)", "toner": "Toner cartridge", "lamp": "Desk lamp"}}, | |
| "total_ok": {"type": "noul", "instructions": "Does the TOTAL equal the subtotal plus tax?"}, | |
| "vendor": {"type": "choice", "instructions": "Who issued the document?", | |
| "criteria": {"acme": "ACME Office Supplies", "northwind": "Northwind Ltd", "unknown": "Not shown"}}, | |
| } | |
| def _j(x): | |
| return json.dumps(x, indent=2, ensure_ascii=False) | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .sd-wrap { display: flex; flex-direction: column; gap: 12px; } | |
| .sd-card { border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 10px 14px; } | |
| .sd-head { display: flex; gap: 10px; align-items: baseline; flex-wrap: wrap; } | |
| .sd-type { font-size: 0.8em; opacity: 0.7; border: 1px solid var(--border-color-primary); border-radius: 6px; padding: 0 6px; } | |
| .sd-hl { margin-left: auto; font-weight: 600; } | |
| .sd-q { font-size: 0.9em; opacity: 0.8; margin: 4px 0 8px; } | |
| .sd-row { display: grid; grid-template-columns: minmax(80px, 30%) 1fr 48px; gap: 8px; align-items: center; margin: 3px 0; } | |
| .sd-name { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 0.92em; } | |
| .sd-best { font-weight: 700; } | |
| .sd-bar { background: var(--background-fill-secondary); border-radius: 4px; height: 12px; overflow: hidden; display: block; } | |
| .sd-bar > span { display: block; height: 100%; background: var(--color-accent); border-radius: 4px; } | |
| .sd-p { text-align: right; font-variant-numeric: tabular-nums; font-size: 0.9em; } | |
| """ | |
| with gr.Blocks(title="Super Decider") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# Super Decider\n" | |
| "Send a **state** (text, JSON or images) and **typed questions**; get one calibrated answer per question: " | |
| "a `choice` with a probability per option, a `noul` (P(yes)), or a `score` over an ordered rubric. " | |
| "All answers are read in one denoising step of a fine-tuned DiffusionGemma 26B-A4B (NVFP4).\n\n" | |
| "[Model](https://huggingface.co/Singularity-Compute/super-decider) · " | |
| "[GitHub](https://github.com/Singularity-Compute/super-decider) · " | |
| "[API format](https://github.com/Singularity-Compute/super-decider/blob/main/docs/API.md)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| state = gr.Textbox(label="State (text or JSON)", lines=6, value=EX_TICKET_STATE) | |
| questions = gr.Code(label="Questions (Jev / System One JSON)", language="json", | |
| value=_j(EX_TICKET_Q), lines=14) | |
| thinking = gr.Radio(["off", "auto", "always"], value="auto", label="Thinking", | |
| info="auto: think only when unsure (text, ≤2 questions); always: think first") | |
| images = gr.Gallery(label="Images (optional, up to 8)", type="filepath", interactive=True, | |
| columns=4, height=180) | |
| run = gr.Button("Decide", variant="primary") | |
| with gr.Accordion("Advanced", open=False): | |
| seed = gr.Number(label="Seed (canvas noise)", value=42, precision=0) | |
| with gr.Column(scale=1): | |
| out_html = gr.HTML(label="Answers") | |
| out_json = gr.JSON(label="Response (Jev shape)") | |
| run.click(decide, inputs=[state, questions, thinking, images, seed], outputs=[out_html, out_json], | |
| api_name="decide") | |
| gr.Examples( | |
| label="Text examples", | |
| examples=[ | |
| [EX_TICKET_STATE, _j(EX_TICKET_Q), "auto"], | |
| [EX_INVOICE_STATE, _j(EX_INVOICE_Q), "always"], | |
| [_j(EX_ROUTE_STATE), _j(EX_ROUTE_Q), "auto"], | |
| ], | |
| inputs=[state, questions, thinking], | |
| outputs=[out_html, out_json], | |
| fn=decide, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| gr.Examples( | |
| label="Image example", | |
| examples=[[EX_IMAGE_STATE, _j(EX_IMAGE_Q), "off", ["examples/invoice-demo.png"]]], | |
| inputs=[state, questions, thinking, images], | |
| outputs=[out_html, out_json], | |
| fn=decide, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |