| """MLX Jeeves runner: no-thinking decisions and optional ordinary greedy reasoning. |
| |
| This is a correctness-oriented runner, not a port of the CUDA speculative server. |
| The pointer head stays FP32. Each question gets its own recurrent/KV cache. |
| """ |
| from __future__ import annotations |
| import argparse |
| import json |
| import platform |
| import sys |
| import time |
| from pathlib import Path |
| import numpy as np |
| import mlx.core as mx |
| from mlx_lm import load |
|
|
| |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from dataformat import DataFormat |
| from jeeves_format import Encoder, readout_positions |
|
|
|
|
| def head_scores(weights, hidden, temperature): |
| if not np.isfinite(temperature) or temperature <= 0: |
| raise ValueError("The exported calibration temperature must be finite and positive.") |
| h = mx.array(hidden, dtype=mx.float32) |
| options = h[:-1] @ weights["k.weight"].T + weights["k.bias"] |
| query = h[-1] @ weights["q.weight"].T + weights["q.bias"] |
| raw = (options @ query) * (weights["q.weight"].shape[0] ** -0.5) |
| calibrated = raw / temperature |
| probabilities = mx.softmax(calibrated, axis=-1) |
| mx.eval(raw, calibrated, probabilities) |
| result = tuple(np.asarray(x, dtype=np.float32) for x in (raw, calibrated, probabilities)) |
| if not all(np.isfinite(x).all() for x in result) or not np.isclose(result[2].sum(), 1.0, atol=1e-5): |
| raise FloatingPointError("Non-finite or invalid decision probabilities.") |
| return result |
|
|
|
|
| class JeevesMLX: |
| def __init__(self, model_dir, *, device="auto", chunk_size=16, max_context=4096): |
| self.path = Path(model_dir).expanduser().resolve() |
| if device == "auto": |
| device = "gpu" if platform.system() == "Darwin" and mx.metal.is_available() else "cpu" |
| if device not in {"cpu", "gpu"}: |
| raise ValueError("device must be auto, cpu, or gpu") |
| if device == "gpu" and not mx.metal.is_available(): |
| raise RuntimeError("This runner's GPU mode requires Apple Metal. Use --device cpu on Colab.") |
| mx.set_default_device(mx.cpu if device == "cpu" else mx.gpu) |
| self.device = device |
| self.chunk_size = int(chunk_size) |
| self.max_context = int(max_context) |
| if self.chunk_size < 1 or self.max_context < 1: |
| raise ValueError("chunk_size and max_context must be positive") |
| self.export = json.loads((self.path / "export.json").read_text()) |
| self.temperature = float(self.export["temperature"]) |
| self.encoder = Encoder(str(self.path), pad_multiple=1) |
| if self.encoder.FORMAT != self.export["format"]: |
| raise ValueError("The exported prompt format does not match the bundled encoder.") |
| self.model, _ = load(str(self.path), tokenizer_config={"trust_remote_code": False}, lazy=False) |
| self.lm = self.model.language_model |
| self.core = self.lm.model |
| self.model.eval() |
| self.head = {k: v.astype(mx.float32) for k, v in mx.load(str(self.path / "head.safetensors")).items()} |
| mx.eval(self.head) |
| if set(self.head) != {"q.weight", "q.bias", "k.weight", "k.bias"}: |
| raise ValueError("Unexpected pointer-head tensors") |
| self.max_context = min(self.max_context, self.lm.args.max_position_embeddings) |
|
|
| def _evaluate(self, tokens, cache): |
| h = self.core(mx.array(tokens, dtype=mx.int32)[None], cache=cache) |
| mx.eval(h, [entry.state for entry in cache]) |
| return h |
|
|
| def score_tokens(self, ids, n_options, *, progress=False): |
| """Return selected final-normalized hidden states plus head probabilities.""" |
| if not ids or len(ids) > self.max_context: |
| raise ValueError(f"Input length {len(ids)} exceeds cap {self.max_context}; no silent truncation.") |
| if ids[-1] != self.encoder.decide_id: |
| raise ValueError("The last token is not the Jeeves decision marker.") |
| opts, decide = readout_positions(ids, self.encoder.markers) |
| if len(opts) != n_options: |
| raise ValueError("Suffix option marker count mismatch") |
| positions = opts + [decide] |
| selected = {} |
| cache = self.lm.make_cache() |
| for start in range(0, len(ids), self.chunk_size): |
| end = min(start + self.chunk_size, len(ids)) |
| h = self._evaluate(ids[start:end], cache) |
| for pos in positions: |
| if start <= pos < end: |
| selected[pos] = np.asarray(h[0, pos - start].astype(mx.float32)).copy() |
| if progress: |
| print(f" CPU prefill {end}/{len(ids)} tokens", flush=True) |
| hidden = np.stack([selected[pos] for pos in positions]) |
| if not np.isfinite(hidden).all(): |
| raise FloatingPointError("Non-finite hidden states") |
| raw, logits, probabilities = head_scores(self.head, hidden, self.temperature) |
| return {"hidden": hidden, "raw_logits": raw, "logits": logits, "probabilities": probabilities, |
| "option_positions": opts, "decide_position": decide} |
|
|
| def greedy_chain(self, prompt, budget): |
| """No speculative drafting. Stop before </think>; caller adds the suffix.""" |
| if budget < 1 or len(prompt) + budget > self.max_context: |
| raise ValueError("Invalid reasoning budget or context overflow") |
| cache = self.lm.make_cache() |
| for start in range(0, len(prompt), self.chunk_size): |
| h = self._evaluate(prompt[start:start + self.chunk_size], cache) |
| banned = mx.array(self.encoder.banned_ids, dtype=mx.int32) |
| bias = mx.zeros((self.lm.args.vocab_size,), dtype=mx.float32) |
| bias[banned] = -1e4 |
| chain = [] |
| for _ in range(budget): |
| logits = (self.lm.model.embed_tokens.as_linear(h[:, -1]) if self.lm.args.tie_word_embeddings |
| else self.lm.lm_head(h[:, -1])) |
| logits = logits.astype(mx.float32) + bias |
| token = int(mx.argmax(logits[0]).item()) |
| if token == self.encoder.think_end_id: |
| return chain, True |
| chain.append(token) |
| h = self._evaluate([token], cache) |
| return chain, False |
|
|
| def predict(self, request, *, think_tokens=0): |
| if not isinstance(think_tokens, int) or isinstance(think_tokens, bool) or think_tokens < 0: |
| raise ValueError("think_tokens must be a nonnegative integer") |
| record = DataFormat.from_dict(request) |
| if len(record.questions) > 64: |
| raise ValueError("At most 64 questions per request in this reference runner") |
| results = {} |
| for question in record.questions: |
| keys = question.keys() |
| if not keys: |
| raise ValueError(f"Question {question.id} has no options") |
| prompt = self.encoder.encode(self.encoder.prompt_text(record, question)) |
| suffix = self.encoder.encode(self.encoder.remainder_text(question)) |
| if len(prompt) + max(think_tokens, len(self.encoder.empty_think)) + len(suffix) > self.max_context: |
| raise ValueError("Request plus reasoning/suffix exceeds max_context; shorten the request.") |
| started = time.perf_counter() |
| if think_tokens: |
| chain, closed = self.greedy_chain(prompt, think_tokens) |
| else: |
| chain, closed = list(self.encoder.empty_think), True |
| |
| scored = self.score_tokens(prompt + chain + suffix, len(keys)) |
| p = scored["probabilities"] |
| results[question.id] = { |
| "type": question.type, "selected": keys[int(p.argmax())], |
| "probabilities": {key: float(value) for key, value in zip(keys, p)}, |
| "input_tokens": len(prompt) + len(chain) + len(suffix), |
| "reasoning_tokens": len(chain) if think_tokens else 0, |
| "reasoning_closed": closed if think_tokens else None, |
| "reasoning": self.encoder.tok.decode(chain, skip_special_tokens=False) if think_tokens else None, |
| "elapsed_seconds": time.perf_counter() - started, |
| } |
| return {"model": "PostHog/jeeves (MLX port)", "device": self.device, |
| "mode": "greedy-thinking" if think_tokens else "no-thinking", |
| "temperature": self.temperature, "results": results} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", default=str(Path(__file__).resolve().parent)) |
| parser.add_argument("--request", required=True, help="Path to a JSON state/questions object") |
| parser.add_argument("--output", default="") |
| parser.add_argument("--device", choices=["auto", "cpu", "gpu"], default="auto") |
| parser.add_argument("--chunk-size", type=int, default=16) |
| parser.add_argument("--max-context", type=int, default=4096) |
| parser.add_argument("--think-tokens", type=int, default=0) |
| args = parser.parse_args() |
| model = JeevesMLX(args.model, device=args.device, chunk_size=args.chunk_size, max_context=args.max_context) |
| request = json.loads(Path(args.request).read_text(encoding="utf-8")) |
| result = model.predict(request, think_tokens=args.think_tokens) |
| text = json.dumps(result, indent=2, ensure_ascii=False, allow_nan=False) |
| if args.output: |
| Path(args.output).write_text(text + "\n", encoding="utf-8") |
| print(text) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
