| """System One decisions on Apple Silicon: one forward pass, typed answers. |
| |
| A faithful MLX port of the transformers System One runner that ships with |
| LiquidAI/d1-3B (`prompt.py`, `runner.py`, `api.py`). The prompt rendering and |
| the logit readout are identical; the forward pass runs through mlx-vlm. |
| |
| python system_one.py --model DJLougen/d1-3B-MLX-bf16 \ |
| --state '{"text": "I was charged twice this month, please refund one of them."}' \ |
| --questions '{"refund": {"type": "noul", "instructions": "Is this a refund request?"}}' |
| |
| The model sees the state and the question once, and the answer is read off the |
| logits at the answer slot. Nothing is decoded, so a schema violation is |
| impossible and tokens generated per decision is zero. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import sys |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Mapping, Sequence |
|
|
| |
| |
| |
|
|
|
|
| @dataclass(frozen=True) |
| class Choice: |
| instructions: str |
| criteria: Mapping[str, str | None] |
| type: str = "choice" |
|
|
|
|
| @dataclass(frozen=True) |
| class Noul: |
| instructions: str |
| criteria: Mapping[str, Any] | None = None |
| type: str = "noul" |
|
|
|
|
| @dataclass(frozen=True) |
| class Score: |
| instructions: str |
| criteria: Sequence[str] |
| type: str = "score" |
|
|
|
|
| Question = Choice | Noul | Score |
|
|
|
|
| def as_question(q: Mapping | Question) -> Question: |
| """`{"type": "noul" | "choice" | "score", "instructions", "criteria"}` as a class; a class passes through.""" |
| if not isinstance(q, Mapping): |
| return q |
| kind = q.get("type", "choice") |
| if kind == "noul": |
| return Noul(q["instructions"], q.get("criteria")) |
| if kind == "score": |
| return Score(q["instructions"], list(q["criteria"])) |
| return Choice(q["instructions"], q["criteria"]) |
|
|
|
|
| |
| |
| |
|
|
|
|
| def option_codes(labels: Sequence[str]) -> list[str]: |
| """Native letters when the labels already are letters, else A..Z, else 00..""" |
| labs = [str(x).strip() for x in labels] |
| if labs and all(len(k) == 1 and k.isalpha() for k in labs): |
| return labs |
| if len(labs) <= 26: |
| return [chr(ord("A") + i) for i in range(len(labs))] |
| return [f"{i:02d}" for i in range(len(labs))] |
|
|
|
|
| _FALLBACK_POOL = ( |
| [chr(c) for c in range(ord("A"), ord("Z") + 1)] |
| + [f"{i:02d}" for i in range(100)] |
| + [chr(c) for c in range(ord("a"), ord("z") + 1)] |
| + [f"#{i}" for i in range(200)] |
| + [chr(a) + chr(b) for a in range(ord("A"), ord("Z") + 1) for b in range(ord("A"), ord("Z") + 1)] |
| ) |
|
|
|
|
| def aliases(tokenizer, labels: Sequence[str]) -> list[tuple[str, int]]: |
| """Assign every label a distinct single-token code: [(code, token_id)].""" |
| used: set[int] = set() |
| out: list[tuple[str, int]] = [] |
|
|
| def take(raw: str) -> bool: |
| enc = tokenizer.encode(raw, add_special_tokens=False) |
| if len(enc) != 1 or enc[0] in used: |
| return False |
| out.append((raw, enc[0])) |
| used.add(enc[0]) |
| return True |
|
|
| for code in option_codes(labels): |
| if take(code): |
| continue |
| if not any(take(raw) for raw in _FALLBACK_POOL): |
| raise RuntimeError(f"no single-token alias left for {len(option_codes(labels))} options") |
| return out |
|
|
|
|
| def _ids(tokenizer, texts: Sequence[str]) -> list[int]: |
| """Single-token ids for the texts that encode to one token, deduped, in order.""" |
| out, seen = [], set() |
| for t in texts: |
| enc = tokenizer.encode(t, add_special_tokens=False) |
| if len(enc) == 1 and enc[0] not in seen: |
| out.append(enc[0]) |
| seen.add(enc[0]) |
| return out |
|
|
|
|
| YES_FORMS = ("yes", "Yes", "YES") |
| NO_FORMS = ("no", "No", "NO") |
|
|
|
|
| def readout_ids(tokenizer, q: Question) -> list[list[int]]: |
| """Token ids to score, one group per option, max-pooled.""" |
| if isinstance(q, Noul): |
| yes, no = _ids(tokenizer, YES_FORMS), _ids(tokenizer, NO_FORMS) |
| if not yes or not no: |
| raise RuntimeError("tokenizer has no single-token yes/no") |
| return [yes, no] |
| if isinstance(q, Score): |
| groups = [_ids(tokenizer, [str(i)]) for i in range(len(q.criteria))] |
| if any(not g for g in groups): |
| raise RuntimeError( |
| f"score with {len(q.criteria)} levels needs single-token digits; " |
| "the primitive is defined for 2 to 10" |
| ) |
| return groups |
| groups = [] |
| for code, tid in aliases(tokenizer, list(q.criteria.keys())): |
| extra = _ids(tokenizer, [f" {code}"]) |
| groups.append([tid] + [i for i in extra if i != tid]) |
| if not groups: |
| raise RuntimeError("choice with no options") |
| return groups |
|
|
|
|
| def readout(tokenizer, q: Question, logz, calibration=None) -> list[float]: |
| """Option probabilities from the log-probabilities at the answer slot. |
| |
| `logz` is indexed by token id (a 1-D array over the vocabulary). Each |
| option scores its best form; the answer is a softmax over these and |
| nothing else. |
| """ |
| scores = [max(float(logz[i]) for i in g) for g in readout_ids(tokenizer, q)] |
| if calibration is not None: |
| scores = calibration.apply(q, scores) |
| m = max(scores) |
| exps = [math.exp(s - m) for s in scores] |
| return [e / sum(exps) for e in exps] |
|
|
|
|
| |
| |
| |
|
|
| IM_START = "<|im_start|>" |
| IM_END = "<|im_end|>" |
|
|
|
|
| def state_block(state: Any) -> str: |
| """The default `json_only` style: a string passes through, anything else as JSON.""" |
| if isinstance(state, str): |
| return f"{state}\n\n" |
| return json.dumps(state, ensure_ascii=False, indent=2) + "\n\n" |
|
|
|
|
| def _option_line(code: str, label: str, desc: str | None) -> str: |
| return f"{code} {desc or label.replace('_', ' ')}" |
|
|
|
|
| def question_block(tokenizer, q: Question) -> str: |
| if isinstance(q, Choice): |
| labels = list(q.criteria.keys()) |
| codes = aliases(tokenizer, labels) |
| lines = "\n".join(_option_line(codes[i][0], lab, q.criteria[lab]) for i, lab in enumerate(labels)) |
| return f"{q.instructions}\n\nOptions:\n{lines}\n\nReply with the option code only." |
| if isinstance(q, Noul): |
| extra = "" |
| if q.criteria: |
| extra = f"\nYes: {q.criteria.get('true')}\nNo: {q.criteria.get('false')}" |
| return f"{q.instructions}{extra}\n\nReply with yes or no only." |
| if isinstance(q, Score): |
| legend = "\n".join(f"{i} {name}" for i, name in enumerate(q.criteria)) |
| return f"{q.instructions}\n\n{legend}\n\nReply with a single digit 0-{len(q.criteria) - 1} only." |
| raise TypeError(f"unknown question type {type(q)}") |
|
|
|
|
| def prefix_text(state: Any, bos: str = "", images: str = "") -> str: |
| """Everything before the question: pictures' markup at the head of the user turn, then the state.""" |
| body = "" if state is None else f"{state_block(state)}\nQUESTION:\n" |
| return f"{bos}{IM_START}user\n{images}{body}" |
|
|
|
|
| def suffix_text(tokenizer, q: Question) -> str: |
| """The question and the assistant header, up to the answer slot (no lead on LFM2-VL).""" |
| return f"{question_block(tokenizer, q)}{IM_END}\n{IM_START}assistant\n" |
|
|
|
|
| def render(tokenizer, state: Any, q: Question, bos: str = "", images: str = "") -> str: |
| return prefix_text(state, bos, images) + suffix_text(tokenizer, q) |
|
|
|
|
| def answer(q: Question, probs: Sequence[float]) -> dict: |
| """A noul's P(yes); a choice's pick and its probabilities; a score's expected level.""" |
| if isinstance(q, Noul): |
| return {"type": "noul", "noul": probs[0]} |
| best = max(range(len(probs)), key=probs.__getitem__) |
| if isinstance(q, Choice): |
| names = list(q.criteria) |
| return {"type": "choice", "choice": names[best], "confidence": probs[best], |
| "probabilities": dict(zip(names, probs))} |
| return {"type": "score", "score": sum(i * p for i, p in enumerate(probs)), "confidence": probs[best], |
| "probabilities": {str(i): p for i, p in enumerate(probs)}, |
| "legend": {str(i): text for i, text in enumerate(q.criteria)}} |
|
|
|
|
| |
| |
| |
|
|
| VISION_MAX_PIXELS = 1024 * 1024 |
|
|
|
|
| def cap_pixels(image, max_pixels: int = VISION_MAX_PIXELS): |
| """A picture downscaled to at most `max_pixels`, bicubic (as upstream does).""" |
| image = image.convert("RGB") if hasattr(image, "convert") else image |
| w, h = image.size |
| if w * h <= max_pixels: |
| return image |
| from PIL import Image |
|
|
| scale = math.sqrt(max_pixels / (w * h)) |
| return image.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.Resampling.BICUBIC) |
|
|
|
|
| class SystemOne: |
| """State in, calibrated distribution out, one forward pass — on MLX.""" |
|
|
| def __init__(self, model_path: str, calibration=None): |
| from mlx_vlm import load |
| from mlx_vlm.utils import load_image |
|
|
| self.model, self.processor = load(model_path) |
| self.tokenizer = self.processor.tokenizer if hasattr(self.processor, "tokenizer") else self.processor |
| self.model_path = model_path |
| self.calibration = calibration |
| bos = getattr(self.tokenizer, "bos_token", None) |
| self.bos = bos if isinstance(bos, str) else "" |
|
|
| |
|
|
| def _image_markup(self, n: int) -> str: |
| """What the chat template writes for `n` images at the head of a user turn.""" |
| msgs = [{"role": "user", "content": [*([{"type": "image"}] * n), {"type": "text", "text": "\x00"}]}] |
| text = self.processor.apply_chat_template(msgs, add_generation_prompt=False, tokenize=False) |
| head = f"{IM_START}user\n" |
| return text[text.index(head) + len(head):text.index("\x00")] |
|
|
| |
|
|
| def _logz(self, prompt: str, images: Sequence = ()): |
| """Log-softmax at the answer slot for one rendered prompt.""" |
| import mlx.core as mx |
| from mlx_vlm.utils import prepare_inputs |
|
|
| pics = [cap_pixels(im) for im in images] |
| inputs = prepare_inputs( |
| self.processor, |
| images=list(pics) or None, |
| prompts=[prompt], |
| image_token_index=getattr(self.model.config, "image_token_index", None), |
| add_special_tokens=False, |
| ) |
| input_ids = inputs["input_ids"] |
| pixel_values = inputs.get("pixel_values", None) |
| mask = inputs.get("attention_mask", None) |
| data_kwargs = {k: v for k, v in inputs.items() |
| if k not in ("input_ids", "pixel_values", "attention_mask")} |
| output = self.model(input_ids, pixel_values, mask, **data_kwargs) |
| logits = output.logits if hasattr(output, "logits") else output |
| row = logits[0, -1] |
| return row - mx.logsumexp(row, axis=-1), int(input_ids.shape[1]) |
|
|
| |
|
|
| def probabilities(self, state: Any, questions: Sequence, images: Sequence | None = None) -> list[list[float]]: |
| """Each question's distribution over its options (`yes`, `no` for a noul). |
| |
| One forward pass per question: the upstream runner packs several |
| questions of one state into a shared trunk; this helper re-reads the |
| state per question — same answers, more compute. |
| """ |
| qs = [as_question(q) for q in questions] |
| pics = [cap_pixels(im) for im in (images or ())] |
| markup = self._image_markup(len(pics)) if pics else "" |
| out = [] |
| for q in qs: |
| logz, _ = self._logz(render(self.tokenizer, state, q, self.bos, markup), pics) |
| out.append(readout(self.tokenizer, q, logz, self.calibration)) |
| return out |
|
|
| def system_one(self, state: Any, questions: Mapping[str, Any], images: Sequence | None = None) -> dict: |
| """Named questions over one state, and its pictures if any, in one forward pass each.""" |
| qs = {n: as_question(q) for n, q in questions.items()} |
| probs = self.probabilities(state, list(qs.values()), images) |
| return {"answers": {n: answer(q, p) for (n, q), p in zip(qs.items(), probs)}, |
| "usage": {"input_tokens": 0, "output_tokens": 0}} |
|
|
|
|
| |
| |
| |
|
|
|
|
| def main(argv=None): |
| parser = argparse.ArgumentParser(description="System One decisions with MLX") |
| parser.add_argument("--model", required=True, help="MLX model path or Hugging Face ID") |
| parser.add_argument("--state", help="Text to evaluate") |
| parser.add_argument("--state-file", type=Path, help="UTF-8 text file to evaluate") |
| parser.add_argument("--questions", type=json.loads, help="Named questions as JSON") |
| parser.add_argument("--questions-file", type=Path, help="JSON file of named questions") |
| parser.add_argument("--image", action="append", default=[], help="Image path or URL (repeatable)") |
| args = parser.parse_args(argv) |
| try: |
| state = args.state_file.read_text(encoding="utf-8") if args.state_file is not None else args.state |
| questions = (json.loads(args.questions_file.read_text(encoding="utf-8")) |
| if args.questions_file is not None else args.questions) |
| for name, q in questions.items(): |
| as_question(q) |
| except (OSError, ValueError, KeyError, TypeError) as error: |
| parser.error(str(error)) |
| engine = SystemOne(args.model) |
| images = None |
| if args.image: |
| from mlx_vlm.utils import load_image |
|
|
| images = [load_image(p) for p in args.image] |
| result = engine.system_one(state, questions, images) |
| print(json.dumps(result, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
