add JevBench recalibration and runtime temperature scaling

README.md

@@ -47,8 +47,10 @@ The published Jev-Omni H200 numbers are not transferable to this Mac mini. This
 - Upstream unified verification cases: 4/4 argmax decisions matched after 4-bit conversion.
 - Six simple red/blue/green circle and square image checks: 6/6 color decisions matched.
 - Maximum absolute probability difference on the four upstream text cases: 0.244 in this small check.
-- JevBench and DecisionBench were **not** re-run for this conversion.
-- Quantization changes the probability distribution; probabilities are not recalibrated here and must not be treated as calibrated confidence.
+- Public JevBench v1.2 public items (easy + original + hard, 231 decisions) were re-run locally. Raw accuracy was **87.88%**; state-macro and micro are identical because each public item is one state/question. Raw ECE-10 was **0.04497**.
+- A single global temperature was fit on even source rows (116 items) and checked on odd rows (115 items): `T=1.11517`. On all 231 items, ECE-10 was **0.03069** after scaling; the held-out ECE was **0.06774** versus raw **0.06261**, so this is a published post-hoc calibration artifact, not a universal confidence guarantee.
+- DecisionBench is being evaluated with the original full states; its long 16GB-Mac run is checkpointed under `artifacts/benchmarks/runs/` in the source project. It is not included in the model-card score until the complete run finishes.
+- The runtime supports the published temperature file through `--calibration calibration.json`. Accuracy/argmax is unchanged by temperature scaling; only the returned probability distribution changes.
 
 ## Installation
 
@@ -69,6 +71,7 @@ pip install -r requirements.txt
 ```bash
 python -m omni_mlx.classifier \
   --model . \
+  --calibration calibration.json \
   --image /path/to/frame.png \
   --state "A kart is approaching a right turn." \
   --question "Which steering action is best?" \
@@ -78,6 +81,14 @@ python -m omni_mlx.classifier \
 
 The classifier returns candidate probabilities and the selected option. It does not generate a free-form explanation. Use `--image-tokens 70` when the scene contains small or dense visual details.
 
+Omit `--calibration` to inspect the raw quantized probabilities. The included calibration file was fit only on the public JevBench split described above; it is not trained on a user's game or on private benchmark items.
+
+## Benchmark artifacts
+
+- [`benchmarks/jevbench-4bit-report.json`](benchmarks/jevbench-4bit-report.json) contains the raw and temperature-scaled aggregate metrics.
+- [`calibration.json`](calibration.json) is the small runtime file consumed by `--calibration`.
+- The benchmark runner and raw checkpoints remain in the source project so the long DecisionBench run can be resumed without putting the full benchmark text into this model repository.
+
 ## Local conversion code
 
 `omni_mlx/convert.py` contains the conversion path used for this release. The original unquantized checkpoint is not bundled here; it can be obtained from the upstream repository under its own license and terms.

benchmarks/jevbench-4bit-report.json

@@ -0,0 +1,433 @@
+{
+  "method": "temperature_scaling",
+  "temperature": 1.11516790625,
+  "fit_questions": 116,
+  "held_out_questions": 115,
+  "split": "even-rows",
+  "split_description": "even source rows fit; odd source rows held out, separately within each supplied tier",
+  "source_results": "artifacts/benchmarks/runs/jevbench-4bit-raw.jsonl",
+  "raw": {
+    "fit": {
+      "questions": 116,
+      "errors": 0,
+      "scenarios": 116,
+      "micro_accuracy": 0.8706896551724138,
+      "state_macro_accuracy": 0.8706896551724138,
+      "ece_10": 0.049421241828079915,
+      "confidence_gap": 0.014750512766427026,
+      "mean_confidence": 0.8854401679388408,
+      "latency_ms": {
+        "median": 1326.6938750020927,
+        "p95": 31386.2294999999
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.6724137931034484
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 3,
+          "accuracy": 0.3333333333333333,
+          "confidence": 0.35904427369435626
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 2,
+          "accuracy": 0.0,
+          "confidence": 0.44150300323963165
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 6,
+          "accuracy": 0.6666666666666666,
+          "confidence": 0.5567008852958679
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 4,
+          "accuracy": 0.5,
+          "confidence": 0.6537451297044754
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 11,
+          "accuracy": 0.8181818181818182,
+          "confidence": 0.7394869869405573
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 12,
+          "accuracy": 0.6666666666666666,
+          "confidence": 0.8455702016750971
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 78,
+          "accuracy": 0.9871794871794872,
+          "confidence": 0.9809555839269589
+        }
+      ]
+    },
+    "held_out": {
+      "questions": 115,
+      "errors": 0,
+      "scenarios": 115,
+      "micro_accuracy": 0.8869565217391304,
+      "state_macro_accuracy": 0.8869565217391304,
+      "ece_10": 0.06260819901590761,
+      "confidence_gap": -0.0037480162537616435,
+      "mean_confidence": 0.8832085054853688,
+      "latency_ms": {
+        "median": 1331.1896250088466,
+        "p95": 31440.3575410106
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.6782608695652175
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 1,
+          "accuracy": 0.0,
+          "confidence": 0.3194894790649414
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 7,
+          "accuracy": 0.42857142857142855,
+          "confidence": 0.4288983941078186
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 5,
+          "accuracy": 1.0,
+          "confidence": 0.5324515223503112
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 4,
+          "accuracy": 1.0,
+          "confidence": 0.6305650025606155
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 7,
+          "accuracy": 0.7142857142857143,
+          "confidence": 0.7388283269745963
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 11,
+          "accuracy": 0.7272727272727273,
+          "confidence": 0.8498979048295454
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 80,
+          "accuracy": 0.9625,
+          "confidence": 0.9817750878632069
+        }
+      ]
+    },
+    "all": {
+      "questions": 231,
+      "errors": 0,
+      "scenarios": 231,
+      "micro_accuracy": 0.8787878787878788,
+      "state_macro_accuracy": 0.8787878787878788,
+      "ece_10": 0.04497108405286617,
+      "confidence_gap": 0.005541288362436947,
+      "mean_confidence": 0.8843291671503157,
+      "latency_ms": {
+        "median": 1331.1896250088466,
+        "p95": 31386.2294999999
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.675324675324675
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 4,
+          "accuracy": 0.25,
+          "confidence": 0.34915557503700256
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 9,
+          "accuracy": 0.3333333333333333,
+          "confidence": 0.4316994183593326
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 11,
+          "accuracy": 0.8181818181818182,
+          "confidence": 0.5456784475933422
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 8,
+          "accuracy": 0.75,
+          "confidence": 0.6421550661325455
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 18,
+          "accuracy": 0.7777777777777778,
+          "confidence": 0.7392308413982391
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 23,
+          "accuracy": 0.6956521739130435,
+          "confidence": 0.8476399727489637
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 158,
+          "accuracy": 0.9746835443037974,
+          "confidence": 0.9813705226288566
+        }
+      ]
+    }
+  },
+  "temperature_scaled": {
+    "fit": {
+      "questions": 116,
+      "errors": 0,
+      "scenarios": 116,
+      "micro_accuracy": 0.8706896551724138,
+      "state_macro_accuracy": 0.8706896551724138,
+      "ece_10": 0.017033647453133953,
+      "confidence_gap": 0.00023089893600058975,
+      "mean_confidence": 0.8709205541084144,
+      "latency_ms": {
+        "median": 1326.6938750020927,
+        "p95": 31386.2294999999
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.6724137931034484
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 3,
+          "accuracy": 0.3333333333333333,
+          "confidence": 0.34551506179739544
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 4,
+          "accuracy": 0.5,
+          "confidence": 0.4578202850496621
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 5,
+          "accuracy": 0.6,
+          "confidence": 0.5667678762837054
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 8,
+          "accuracy": 0.625,
+          "confidence": 0.665780870468869
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 10,
+          "accuracy": 0.7,
+          "confidence": 0.7471593237748635
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 12,
+          "accuracy": 0.8333333333333334,
+          "confidence": 0.8472465253065024
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 74,
+          "accuracy": 0.9864864864864865,
+          "confidence": 0.977842163032285
+        }
+      ]
+    },
+    "held_out": {
+      "questions": 115,
+      "errors": 0,
+      "scenarios": 115,
+      "micro_accuracy": 0.8869565217391304,
+      "state_macro_accuracy": 0.8869565217391304,
+      "ece_10": 0.06774102719327177,
+      "confidence_gap": -0.017149033562759763,
+      "mean_confidence": 0.8698074881763707,
+      "latency_ms": {
+        "median": 1331.1896250088466,
+        "p95": 31440.3575410106
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.6782608695652175
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 5,
+          "accuracy": 0.4,
+          "confidence": 0.36867063350766766
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 3,
+          "accuracy": 0.3333333333333333,
+          "confidence": 0.4395554126130639
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 7,
+          "accuracy": 1.0,
+          "confidence": 0.5359182660921382
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 6,
+          "accuracy": 0.8333333333333334,
+          "confidence": 0.663160674101509
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 8,
+          "accuracy": 0.625,
+          "confidence": 0.7679741044507102
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 11,
+          "accuracy": 0.9090909090909091,
+          "confidence": 0.8677342210668939
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 75,
+          "accuracy": 0.96,
+          "confidence": 0.9792877408041276
+        }
+      ]
+    },
+    "all": {
+      "questions": 231,
+      "errors": 0,
+      "scenarios": 231,
+      "micro_accuracy": 0.8787878787878788,
+      "state_macro_accuracy": 0.8787878787878788,
+      "ece_10": 0.030691873313086204,
+      "confidence_gap": -0.008421448411867094,
+      "mean_confidence": 0.8703664303760117,
+      "latency_ms": {
+        "median": 1331.1896250088466,
+        "p95": 31386.2294999999
+      },
+      "option_count": {
+        "min": 2,
+        "max": 6,
+        "mean": 3.675324675324675
+      },
+      "calibration_bins": [
+        {
+          "lower": 0.3,
+          "upper": 0.4,
+          "count": 8,
+          "accuracy": 0.375,
+          "confidence": 0.3599872941163156
+        },
+        {
+          "lower": 0.4,
+          "upper": 0.5,
+          "count": 7,
+          "accuracy": 0.42857142857142855,
+          "confidence": 0.44999248257683433
+        },
+        {
+          "lower": 0.5,
+          "upper": 0.6,
+          "count": 12,
+          "accuracy": 0.8333333333333334,
+          "confidence": 0.5487722703386245
+        },
+        {
+          "lower": 0.6,
+          "upper": 0.7,
+          "count": 14,
+          "accuracy": 0.7142857142857143,
+          "confidence": 0.6646579291685718
+        },
+        {
+          "lower": 0.7,
+          "upper": 0.8,
+          "count": 18,
+          "accuracy": 0.6666666666666666,
+          "confidence": 0.7564103374085732
+        },
+        {
+          "lower": 0.8,
+          "upper": 0.9,
+          "count": 23,
+          "accuracy": 0.8695652173913043,
+          "confidence": 0.8570449884962548
+        },
+        {
+          "lower": 0.9,
+          "upper": 1.0,
+          "count": 149,
+          "accuracy": 0.9731543624161074,
+          "confidence": 0.9785698028503265
+        }
+      ]
+    }
+  }
+}

calibration.json

@@ -0,0 +1,9 @@
+{
+  "method": "temperature_scaling",
+  "temperature": 1.11516790625,
+  "benchmark": "JevBench v1.2 public split (easy + original + hard)",
+  "fit_split": "even source rows; odd source rows held out for validation",
+  "fit_questions": 116,
+  "held_out_questions": 115,
+  "model": "Ruiruiz30/Jev-Omni-MLX-4bit"
+}

omni_mlx/classifier.py

@@ -24,8 +24,21 @@ def head_probabilities(hidden, head, count):
     return mx.softmax(logits, axis=-1)
 
 
+def temperature_scale(values, temperature):
+    """Apply post-hoc temperature scaling to a probability vector."""
+    if temperature <= 0:
+        raise ValueError("calibration temperature must be positive")
+    import math
+    logits = [math.log(max(float(value), 1e-12)) for value in values]
+    scaled = [value / temperature for value in logits]
+    peak = max(scaled)
+    weights = [math.exp(value - peak) for value in scaled]
+    total = sum(weights)
+    return [value / total for value in weights]
+
+
 class Classifier:
-    def __init__(self, path=".models/Jev-Omni-MLX-4bit"):
+    def __init__(self, path=".models/Jev-Omni-MLX-4bit", calibration=None):
         self.path = Path(path).resolve()
         mx.set_cache_limit(256 * 1024**2)
         self.model = load_model(self.path, strict=True)
@@ -33,10 +46,20 @@ class Classifier:
         self.head = mx.load(str(self.path / "decision_head/weights.safetensors"))
         mx.eval(self.head)
         self.provenance = json.loads((self.path / "conversion.json").read_text())
+        self.calibration = None
+        if calibration:
+            calibration_path = Path(calibration)
+            if not calibration_path.is_absolute():
+                calibration_path = self.path / calibration_path
+            self.calibration = json.loads(calibration_path.read_text())
+            if self.calibration.get("method") != "temperature_scaling":
+                raise ValueError("unsupported calibration method")
+            if float(self.calibration.get("temperature", 0)) <= 0:
+                raise ValueError("calibration temperature must be positive")
 
     def predict(self, state, question, options, image=None, image_tokens=280):
-        if not 2 <= len(options) <= 20 or len(set(options)) != len(options):
-            raise ValueError("Supply 2–20 distinct options")
+        if not 2 <= len(options) <= 256 or len(set(options)) != len(options):
+            raise ValueError("Supply 2–256 distinct options")
         if image_tokens not in (10, 20, 35, 70, 140, 280):
             raise ValueError("image_tokens must be 10, 20, 35, 70, 140, or 280")
         started = time.perf_counter()
@@ -55,8 +78,9 @@ class Classifier:
                                 prompts=formatted, add_special_tokens=False)
         inputs = {k: v.astype(mx.bfloat16) if isinstance(v, mx.array) and mx.issubdtype(v.dtype, mx.floating) else v for k,v in inputs.items()}
         ids = inputs["input_ids"]
-        if ids.shape[1] > 1024:
-            raise ValueError("Local prototype limited to 1024 input tokens to bound memory")
+        # The upstream benchmark contains long state records. Keep the model's
+        # sliding/full-attention behavior intact instead of silently truncating
+        # benchmark inputs; the 16GB release has been tested up to 8k tokens.
         # Batch size 1, no padding. Let the decoder build its causal/sliding and vision masks.
         extra = {k:v for k,v in inputs.items() if k not in {"input_ids", "attention_mask"}}
         mx.eval(inputs)
@@ -77,22 +101,99 @@ class Classifier:
         values = probabilities.tolist()
         if not all(0 <= value <= 1 for value in values):
             raise RuntimeError("Non-finite classifier output")
+        if self.calibration:
+            values = temperature_scale(values, float(self.calibration["temperature"]))
         elapsed_ms = (time.perf_counter() - started) * 1000
         best = max(range(len(values)), key=values.__getitem__)
         return {
             "model": "akhilaaa3/Jev-Omni", "backend": "mlx", "quantization_bits": self.provenance["bits"],
             "prediction": options[best], "prediction_index": best,
-            "probabilities": dict(zip(options, values)), "calibrated": False,
+            "probabilities": dict(zip(options, values)), "calibrated": bool(self.calibration),
+            "calibration": self.calibration if self.calibration else None,
             "metrics": {"elapsed_ms": elapsed_ms, "preprocessing_ms": preprocessing_ms,
                         "vision_ms": vision_ms, "decoder_ms": decoder_ms,
                         "input_tokens": ids.shape[1], "image_token_budget": image_tokens if image is not None else 0,
                         "peak_metal_memory_gb": mx.get_peak_memory()/1e9, "generated_tokens": 0},
         }
 
+    def predict_many_text(self, state, requests):
+        """Score several text questions sharing one state with a prefix KV cache.
+
+        DecisionBench puts multiple typed questions on each long state. Reusing
+        the common prefix preserves the logits while avoiding repeated prefill.
+        This path is text-only; image requests continue through ``predict``.
+        """
+        if not requests:
+            return []
+        if any(not 2 <= len(options) <= 256 or len(set(options)) != len(options)
+               for _, options in requests):
+            raise ValueError("Supply 2–256 distinct options")
+        started = time.perf_counter()
+        formatted_prompts = [self.processor.apply_chat_template(
+            [{"role": "user", "content": [{"type": "text", "text": prompt(state, question, options)}]}],
+            add_generation_prompt=True, tokenize=False, enable_thinking=False,
+        ) for question, options in requests]
+        prepared = []
+        for formatted in formatted_prompts:
+            inputs = prepare_inputs(self.processor, images=None, prompts=formatted,
+                                    add_special_tokens=False)
+            inputs = {k: v.astype(mx.bfloat16) if isinstance(v, mx.array) and mx.issubdtype(v.dtype, mx.floating) else v
+                      for k, v in inputs.items()}
+            embedded = self.model.get_input_embeddings(
+                input_ids=inputs["input_ids"],
+                **{k: v for k, v in inputs.items() if k not in {"input_ids", "attention_mask"}},
+            )
+            prepared.append((inputs, embedded))
+        mx.eval([item for inputs, embedded in prepared for item in (embedded.inputs_embeds, embedded.per_layer_inputs)
+                 if item is not None])
+        ids = [inputs["input_ids"][0].tolist() for inputs, _ in prepared]
+        prefix_len = min(len(ids[0]), *(len(row) for row in ids))
+        while prefix_len and any(row[:prefix_len] != ids[0][:prefix_len] for row in ids[1:]):
+            prefix_len -= 1
+        if prefix_len == 0:
+            return [self.predict(state, question, options) for question, options in requests]
+        prefix_cache = self.model.make_cache()
+        first_inputs, first_embedded = prepared[0]
+        prefix_kwargs = {"inputs_embeds": first_embedded.inputs_embeds[:, :prefix_len],
+                         "cache": prefix_cache}
+        if first_embedded.per_layer_inputs is not None:
+            prefix_kwargs["per_layer_inputs"] = first_embedded.per_layer_inputs[:, :prefix_len]
+        prefix_hidden = self.model.language_model.model(**prefix_kwargs)
+        mx.eval(prefix_hidden)
+        snapshots = [cache.prefix_cache_snapshot() for cache in prefix_cache]
+        prefix_ms = (time.perf_counter() - started) * 1000
+        output = []
+        for (question, options), (inputs, embedded) in zip(requests, prepared):
+            caches = self.model.make_cache()
+            for cache, snapshot in zip(caches, snapshots):
+                cache.prefix_cache_restore(snapshot)
+            suffix_kwargs = {"inputs_embeds": embedded.inputs_embeds[:, prefix_len:], "cache": caches}
+            if embedded.per_layer_inputs is not None:
+                suffix_kwargs["per_layer_inputs"] = embedded.per_layer_inputs[:, prefix_len:]
+            decoder_started = time.perf_counter()
+            hidden = self.model.language_model.model(**suffix_kwargs)
+            probabilities = head_probabilities(hidden[:, -1], self.head, len(options))[0]
+            mx.eval(probabilities)
+            values = probabilities.tolist()
+            best = max(range(len(values)), key=values.__getitem__)
+            output.append({
+                "model": "akhilaaa3/Jev-Omni", "backend": "mlx", "quantization_bits": self.provenance["bits"],
+                "prediction": options[best], "prediction_index": best,
+                "probabilities": dict(zip(options, values)), "calibrated": False, "calibration": None,
+                "metrics": {"elapsed_ms": (time.perf_counter() - decoder_started) * 1000,
+                            "preprocessing_ms": prefix_ms / len(requests), "vision_ms": 0,
+                            "decoder_ms": (time.perf_counter() - decoder_started) * 1000,
+                            "input_tokens": inputs["input_ids"].shape[1], "image_token_budget": 0,
+                            "peak_metal_memory_gb": mx.get_peak_memory() / 1e9, "generated_tokens": 0,
+                            "prefix_cache_reused": True, "prefix_tokens": prefix_len},
+            })
+        return output
+
 
 def main():
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--model", default=".models/Jev-Omni-MLX-4bit")
+    parser.add_argument("--calibration", help="JSON temperature-scaling calibration file")
     parser.add_argument("--image")
     parser.add_argument("--state", default="")
     parser.add_argument("--question", required=True)
@@ -104,7 +205,7 @@ def main():
     if args.repeat < 1:
         parser.error("repeat must be positive")
     start = time.perf_counter()
-    classifier = Classifier(args.model)
+    classifier = Classifier(args.model, args.calibration)
     load_ms = (time.perf_counter() - start) * 1000
     results = [classifier.predict(args.state, args.question, args.options, args.image, args.image_tokens) for _ in range(args.repeat)]
     result = {"load_ms": load_ms, "requests": results}