jul 0.1.0__py3-none-any.whl

jul/__init__.py

@@ -0,0 +1,26 @@
+"""jul — Juste Un LLM: local typed decisions, with the interface of the TypeSafe (Jev) Python SDK.
+
+Swap the import and existing code keeps working:
+
+    # from typesafe_sdk import TypeSafeClient, Choice, Noul, Score
+    from jul import TypeSafeClient, Choice, Noul, Score
+
+Two extras Jev does not have: a `Context` describing the data to sort, and `client.autotune(...)`, a small
+head trained in seconds on labeled examples while the model itself stays untouched.
+"""
+
+from .client import AsyncTypeSafeClient, TypeSafeClient
+from .context import Context
+from .presets import PRESETS, Preset
+from .tuning import TuningReport
+from .types import (Choice, ChoiceAnswer, Noul, NoulAnswer, NoulCriteria, Score, ScoreAnswer,
+                    SystemOneResponse, Usage)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "TypeSafeClient", "AsyncTypeSafeClient",
+    "Choice", "Noul", "NoulCriteria", "Score",
+    "ChoiceAnswer", "NoulAnswer", "ScoreAnswer", "SystemOneResponse", "Usage",
+    "Context", "TuningReport", "Preset", "PRESETS", "__version__",
+]

jul/assets/presets/minicpm5-2b@torch.json

@@ -0,0 +1,67 @@
+{
+  "name": "minicpm5-2b",
+  "repos": {
+    "torch": "openbmb/MiniCPM5-2B"
+  },
+  "backend": "torch",
+  "formulations": [
+    {
+      "name": "one_word",
+      "template": "This text: \"{state}\" means in one word: \"",
+      "layer": 41
+    },
+    {
+      "name": "question_options",
+      "template": "{instructions}\nPossible answers: {options}.\nText: \"{state}\"\nIn one word, the answer is: \"",
+      "layer": 37
+    }
+  ],
+  "tau": 0.04381,
+  "center": "generic",
+  "one_word": [
+    37,
+    0.03975
+  ],
+  "latency_ms": "~155",
+  "quality": "dev accuracy 0.540 ± 0.035 (n=200); not measured on the Jev bench",
+  "notes": "Fitted by `jul models add` on 2026-09-22.",
+  "calibration": {
+    "layers": {
+      "one_word": 41,
+      "question_options": 37,
+      "one_word_only": 37
+    },
+    "center": "generic",
+    "tau": 0.043805317199998264,
+    "tau_one_word": 0.03975009964560149,
+    "tau_task_center": 0.04132507636642313,
+    "dev_accuracy": 0.54,
+    "dev_accuracy_stderr": 0.035242020373412196,
+    "dev_accuracy_task_center": 0.595,
+    "dev_ece": 0.15921275988221167,
+    "dev_accuracy_by_center": {
+      "options": 0.5361,
+      "generic": 0.55,
+      "none": 0.5506
+    },
+    "dev_accuracy_by_set": {
+      "yahootopics": 0.42,
+      "empathetic": 0.42,
+      "massive": 0.6,
+      "financialphrasebank": 0.72
+    },
+    "n_dev": 200,
+    "backend": "torch",
+    "repo": "openbmb/MiniCPM5-2B",
+    "n_layers": 42,
+    "candidate_layers": [
+      20,
+      41
+    ],
+    "n_generic": 195,
+    "latency_ms_p50": 154.81362499849638,
+    "warnings": [],
+    "date": "2026-09-22",
+    "seconds": 141
+  }
+}

jul/backbone.py

@@ -0,0 +1,215 @@
+"""Backbone: one forward pass, tap hidden states at chosen layers, reuse a cached prompt prefix.
+
+The instruction part of every prompt is identical across calls, so its KV cache is computed once
+and each query only pays for its own tokens. When only intermediate layers are needed, the forward
+stops right after the deepest one (the remaining layers are never computed).
+
+The framework lives behind `Backbone`: `backends/mlx.py` (Apple Silicon) and `backends/torch.py`
+(transformers: CUDA, CPU, MPS). Everything above this module only sees numpy arrays.
+`Backbone(name)` picks the backend from `JUL_BACKEND`, else MLX when available, else torch.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import os
+import platform
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+BACKENDS = ("mlx", "torch")
+
+#: Preset name -> repo per backend. A name missing here is used as the repo itself.
+MODELS: dict[str, dict[str, str]] = {
+    "minicpm5-2b": {"mlx": "openbmb/MiniCPM5-2B-MLX", "torch": "openbmb/MiniCPM5-2B"},
+    "qwen3-0.6b": {"mlx": "mlx-community/Qwen3-0.6B-4bit", "torch": "Qwen/Qwen3-0.6B"},
+    "qwen3-1.7b": {"mlx": "mlx-community/Qwen3-1.7B-4bit", "torch": "Qwen/Qwen3-1.7B"},
+    "qwen3.5-9b": {"mlx": "mlx-community/Qwen3.5-9B-4bit", "torch": "Qwen/Qwen3.5-9B"},
+}
+
+#: Size of a group in `PromptTemplate.run_batch`: rows x longest prompt (cached prefix included).
+#: `JUL_BATCH_TOKENS` / `JUL_BATCH_SIZE` override them, read at each call like the other JUL_* variables.
+BATCH_TOKENS = 16384
+BATCH_SIZE = 64
+
+# Layers tapped during extraction, as fractions of the model depth.
+DEFAULT_LAYER_FRACTIONS = (0.25, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0)
+
+
+def _mlx_available() -> bool:
+    return (platform.system() == "Darwin" and platform.machine() == "arm64"
+            and importlib.util.find_spec("mlx") is not None)
+
+
+def resolve_backend(backend: str | None = None) -> str:
+    backend = backend or os.environ.get("JUL_BACKEND")
+    if backend:
+        if backend not in BACKENDS:
+            raise ValueError(f"Unknown backend {backend!r}. Available: {', '.join(BACKENDS)}")
+        return backend
+    if _mlx_available():
+        return "mlx"
+    if importlib.util.find_spec("torch") is not None:
+        return "torch"
+    raise ImportError("No backend installed: pip install 'jul[mlx]' (Apple Silicon) or 'jul[torch]'")
+
+
+def model_key(name: str, backend: str) -> str:
+    """Identifies vectors computed by a model on a backend (saved centers, heads, calibrations).
+
+    MLX keeps the bare name so that contexts saved before backends existed stay valid.
+    """
+    return name if backend == "mlx" else f"{name}@{backend}"
+
+
+def repo_for(name: str, backend: str) -> str:
+    repos = MODELS.get(name)
+    if repos is None:
+        return name
+    if backend not in repos:
+        raise ValueError(f"{name!r} has no {backend} repo")
+    return repos[backend]
+
+
+class Backbone:
+    """A causal LM read by the vector method. `Backbone(name)` returns the resolved backend's subclass.
+
+    Subclasses set `tokenizer` (a Hugging Face tokenizer, for `encode` and the chat template) and
+    `n_layers`, and implement `forward` and `cache_prefix`.
+    """
+
+    backend: str = ""
+
+    def __new__(cls, name: str, backend: str | None = None, **kwargs):
+        if cls is Backbone:
+            backend = resolve_backend(backend)
+            if backend == "mlx":
+                from .backends.mlx import MLXBackbone as cls
+            else:
+                from .backends.torch import TorchBackbone as cls
+        return super().__new__(cls)
+
+    def __init__(self, name: str, backend: str | None = None):
+        self.name = name
+        self.repo = repo_for(name, self.backend)
+        self.key = model_key(name, self.backend)
+
+    def layer_indices(self, fractions=DEFAULT_LAYER_FRACTIONS) -> list[int]:
+        return sorted({max(0, min(self.n_layers - 1, round(f * self.n_layers) - 1)) for f in fractions})
+
+    def encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text, add_special_tokens=False)
+
+    def forward(self, tokens: list[int], layers=(), logits=False, pool: tuple[int, int] | None = None,
+                prefix=None) -> tuple[dict[int, np.ndarray], np.ndarray | None]:
+        """Run tokens through the model, after `prefix` (from `cache_prefix`) when given.
+
+        Returns ({layer: (2d,) float32 features}, float32 last-token logits or None). Features are the
+        last-token hidden state concatenated with the mean hidden state over positions
+        pool=(start, end) of `tokens`. The prefix is left as it was, ready for the next query.
+        """
+        raise NotImplementedError
+
+    def forward_batch(self, queries: list[list[int]], layers=(), pools: list | None = None,
+                      prefix=None) -> list[dict[int, np.ndarray]]:
+        """The features of `forward` for each query, all after the same `prefix`. A backend may run
+        them in one forward; this default runs them one by one."""
+        pools = pools or [None] * len(queries)
+        return [self.forward(q, layers=layers, pool=p, prefix=prefix)[0] for q, p in zip(queries, pools)]
+
+    def cache_prefix(self, tokens: list[int]):
+        """Run `tokens` once and keep the model state after them, for `forward(prefix=...)`."""
+        raise NotImplementedError
+
+    def last_hidden(self, tokens: list[int], prefix=None) -> np.ndarray:
+        """(len(tokens), d) float32: the last layer's hidden states after the final norm, for every token
+        of `tokens` (run after `prefix` when given). The prefix is left as it was. Used by the pointer
+        method (jul/decision.py)."""
+        raise NotImplementedError
+
+    @property
+    def model_dir(self) -> Path:
+        """Local directory of the weights (downloaded on first access for a Hub repo)."""
+        if Path(self.repo).is_dir():
+            return Path(self.repo)
+        from huggingface_hub import snapshot_download
+        return Path(snapshot_download(self.repo))
+
+
+@dataclass
+class PromptTemplate:
+    """A chat prompt split around the user input: a cached prefix and a per-query suffix."""
+
+    backbone: Backbone
+    prefix_text: str
+    suffix_text: str
+    use_prefix_cache: bool = True
+
+    SENTINEL = "⁣QF_INPUT⁣"
+
+    @classmethod
+    def from_user_message(cls, backbone: Backbone, message: str, **kw) -> "PromptTemplate":
+        """`message` must contain {input}; it is rendered with the model's chat template."""
+        rendered = backbone.tokenizer.apply_chat_template(
+            [{"role": "user", "content": message.replace("{input}", cls.SENTINEL)}],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+        prefix, suffix = rendered.split(cls.SENTINEL)
+        return cls(backbone, prefix, suffix, **kw)
+
+    def __post_init__(self):
+        self.prefix_tokens = self.backbone.encode(self.prefix_text)
+        self._n_suffix = len(self.backbone.encode(self.suffix_text))
+        self._prefix = None
+        if self.use_prefix_cache and self.prefix_tokens:
+            self._prefix = self.backbone.cache_prefix(self.prefix_tokens)
+
+    def run(self, text: str, layers=(), logits=False):
+        query = self.backbone.encode(text + self.suffix_text)
+        n_input = max(1, len(query) - self._n_suffix)
+        if self._prefix is None:
+            p = len(self.prefix_tokens)
+            return self.backbone.forward(self.prefix_tokens + query, layers=layers, logits=logits, pool=(p, p + n_input))
+        return self.backbone.forward(query, layers=layers, logits=logits, pool=(0, n_input), prefix=self._prefix)
+
+    def run_batch(self, texts: list[str], layers=()) -> list[dict[int, np.ndarray]]:
+        """The features of `run` for each text. Texts are sorted by length and grouped so that a group
+        holds at most BATCH_TOKENS tokens (rows x longest prompt) and BATCH_SIZE rows."""
+        queries = [self.backbone.encode(t + self.suffix_text) for t in texts]
+        n_inputs = [max(1, len(q) - self._n_suffix) for q in queries]
+        p = len(self.prefix_tokens)
+        if self._prefix is None:
+            seqs, pools = [self.prefix_tokens + q for q in queries], [(p, p + n) for n in n_inputs]
+        else:
+            seqs, pools = queries, [(0, n) for n in n_inputs]
+        max_tokens = int(os.environ.get("JUL_BATCH_TOKENS") or BATCH_TOKENS)
+        max_rows = int(os.environ.get("JUL_BATCH_SIZE") or BATCH_SIZE)
+        out: list = [None] * len(seqs)
+        group: list[int] = []
+
+        def flush():
+            got = self.backbone.forward_batch([seqs[i] for i in group], layers=layers,
+                                              pools=[pools[i] for i in group], prefix=self._prefix)
+            for i, features in zip(group, got):
+                out[i] = features
+
+        for i in sorted(range(len(seqs)), key=lambda i: len(seqs[i])):
+            longest = len(seqs[i]) + (p if self._prefix is not None else 0)
+            if group and ((len(group) + 1) * longest > max_tokens or len(group) >= max_rows):
+                flush()
+                group = []
+            group.append(i)
+        if group:
+            flush()
+        return out
+
+
+def timed(fn, *args, **kwargs):
+    t = time.perf_counter()
+    out = fn(*args, **kwargs)
+    return out, time.perf_counter() - t

jul/backends/__init__.py

@@ -0,0 +1 @@
+"""Framework-specific implementations of `jul.backbone.Backbone`. Imported lazily by `Backbone(name)`."""

jul/backends/mlx.py

@@ -0,0 +1,194 @@
+"""MLX backend (Apple Silicon), through mlx-lm.
+
+`forward_batch` runs several queries in one forward, as the torch backend does: padded on the right,
+under the causal mask alone, since a real token never sees the padding after it.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from mlx_lm import load
+from mlx_lm.models.cache import KVCache, can_trim_prompt_cache, make_prompt_cache
+
+from ..backbone import Backbone
+
+
+class _StopForward(Exception):
+    pass
+
+
+def _rope_fix(repo: str) -> dict | None:
+    """transformers 5 writes the RoPE base under `rope_parameters`; mlx-lm reads `rope_theta` and
+    silently falls back to 10000 when it is missing, which gives wrong answers without any error
+    (JOURNAL §9 duovicies). Pass the right value when a converted config has only the new key."""
+    path = Path(repo) / "config.json"
+    if not path.exists():
+        # a Hub repo: the cached config, else fetch that one file (1 KB, before the weights)
+        from huggingface_hub import hf_hub_download, try_to_load_from_cache
+        cached = try_to_load_from_cache(repo, "config.json")
+        try:
+            path = Path(cached if isinstance(cached, str) else hf_hub_download(repo, "config.json"))
+        except Exception:
+            return None
+    config = json.loads(path.read_text())
+    theta = (config.get("rope_parameters") or {}).get("rope_theta")
+    return {"rope_theta": theta} if theta and "rope_theta" not in config else None
+
+
+class _Tap:
+    """Wraps a transformer block to record its features and optionally halt the forward."""
+
+    def __init__(self, block, idx: int, backbone: "MLXBackbone"):
+        self.block = block
+        self.idx = idx
+        self.bb = backbone
+
+    def __call__(self, *args, **kwargs):
+        h = self.block(*args, **kwargs)
+        if self.idx in self.bb._want:
+            pool = self.bb._pool
+            last = h[mx.arange(h.shape[0]), self.bb._last]
+            # mean in float32, rounded back to the model dtype like `h[:, start:end].mean(1)` was
+            mean = (mx.where(pool, h, 0).astype(mx.float32).sum(1) / pool.sum(1)).astype(h.dtype)
+            # [last token ; mean over the input tokens]
+            self.bb._captured[self.idx] = mx.concatenate([last, mean], axis=-1)
+        if self.bb._stop_at == self.idx:
+            raise _StopForward
+        return h
+
+    def __getattr__(self, name):
+        return getattr(self.block, name)
+
+
+@dataclass
+class _Prefix:
+    n: int
+    cache: list | None = None      # attention KV cache: trimmed back to the prefix after each query
+    snapshot: list | None = None   # recurrent state: copied, each query starts from the copy
+
+
+class MLXBackbone(Backbone):
+    backend = "mlx"
+
+    def __init__(self, name: str, backend: str | None = None):
+        super().__init__(name)
+        self.model, self.tokenizer = load(self.repo, model_config=_rope_fix(self.repo))
+        layers = self.model.layers
+        for i, block in enumerate(layers):
+            layers[i] = _Tap(block, i, self)
+        self.n_layers = len(layers)
+        self._want: set[int] = set()
+        self._stop_at: int | None = None
+        self._captured: dict[int, mx.array] = {}
+        self._last: mx.array | None = None   # (B,) index of each row's last real token
+        self._pool: mx.array | None = None   # (B, T, 1) bool mask of each row's pooled positions
+        pad = getattr(self.tokenizer, "pad_token_id", None)
+        self._pad = pad if pad is not None else (self.tokenizer.eos_token_id or 0)
+        lm = getattr(self.model, "language_model", self.model)  # multimodal wrappers (e.g. Qwen3.5) nest the text model
+        self._inner = lm.model
+        self._lm_head = lm.lm_head if hasattr(lm, "lm_head") else self._inner.embed_tokens.as_linear
+
+    def forward(self, tokens, layers=(), logits=False, pool=None, prefix: _Prefix | None = None, cache=None):
+        """`cache` (an mlx-lm prompt cache, advanced in place) is kept for the dev scripts."""
+        if prefix is None:
+            return self._first(self._run([tokens], cache, layers, logits, [pool]))
+        if prefix.snapshot is not None:
+            return self._first(self._run([tokens], self._restored(prefix.snapshot), layers, logits, [pool]))
+        try:
+            return self._first(self._run([tokens], prefix.cache, layers, logits, [pool]))
+        finally:
+            for c in prefix.cache:
+                if c.offset > prefix.n:
+                    c.trim(c.offset - prefix.n)
+
+    def forward_batch(self, queries, layers=(), pools=None, prefix: _Prefix | None = None):
+        pools = pools or [None] * len(queries)
+        if len(queries) == 1 or (prefix is not None and not _repeatable(prefix)):
+            # a recurrent state (or a rotating window) is not repeated over a batch: one query at a time
+            return super().forward_batch(queries, layers, pools, prefix)
+        cache = None
+        if prefix is not None:
+            # a fresh cache holding the prefix once per row; the template's own cache is not touched
+            cache = make_prompt_cache(self.model)
+            for new, c in zip(cache, prefix.cache):
+                keys, values = c.state
+                new.state = (mx.repeat(keys, len(queries), axis=0), mx.repeat(values, len(queries), axis=0))
+        captured, _ = self._run(queries, cache, layers, False, pools)
+        # every group has its own shape (rows x width): MLX would keep the freed buffers of each one
+        mx.clear_cache()
+        return [{k: v[i] for k, v in captured.items()} for i in range(len(queries))]
+
+    def cache_prefix(self, tokens) -> _Prefix:
+        cache = make_prompt_cache(self.model)
+        self._run([tokens], cache, logits=True)
+        mx.eval([c.state for c in cache])
+        if can_trim_prompt_cache(cache):
+            return _Prefix(len(tokens), cache=cache)
+        # Recurrent state (e.g. Gated DeltaNet in Qwen3.5) cannot be trimmed: keep a copy of the state
+        # after the prefix. Layers replace their state arrays rather than writing into them, so the
+        # copy is never modified.
+        return _Prefix(len(tokens), snapshot=[tuple(c.state) for c in cache])
+
+    def last_hidden(self, tokens, prefix: _Prefix | None = None) -> np.ndarray:
+        self._want, self._stop_at = set(), None
+        cache = None
+        if prefix is not None:
+            cache = self._restored(prefix.snapshot) if prefix.snapshot is not None else prefix.cache
+        try:
+            h = self._inner(mx.array(tokens)[None], cache=cache)[0].astype(mx.float32)
+            mx.eval(h)
+            return np.array(h)
+        finally:
+            if prefix is not None and prefix.cache is not None:
+                for c in prefix.cache:
+                    if c.offset > prefix.n:
+                        c.trim(c.offset - prefix.n)
+
+    def _restored(self, snapshot):
+        cache = make_prompt_cache(self.model)
+        for c, state in zip(cache, snapshot):
+            c.state = list(state)
+        return cache
+
+    @staticmethod
+    def _first(result):
+        captured, logits = result
+        return {k: v[0] for k, v in captured.items()}, (logits[0] if logits is not None else None)
+
+    def _run(self, seqs, cache=None, layers=(), logits=False, pools=None):
+        """Right-padded batch of token lists. Returns ({layer: (B, 2d)}, (B, vocab) or None)."""
+        lengths = [len(s) for s in seqs]
+        width = max(lengths)
+        ids = np.full((len(seqs), width), self._pad, dtype=np.int32)
+        pool = np.zeros((len(seqs), width, 1), dtype=bool)
+        for i, (s, p) in enumerate(zip(seqs, pools or [None] * len(seqs))):
+            ids[i, : len(s)] = s
+            start, end = p or (0, len(s))
+            pool[i, start:end] = True
+        self._want = set(layers)
+        self._last = mx.array([n - 1 for n in lengths])
+        self._pool = mx.array(pool)
+        self._stop_at = None if logits else (max(layers) if layers else None)
+        self._captured = {}
+        out = None
+        try:
+            h = self._inner(mx.array(ids), cache=cache)
+            if logits:
+                out = self._lm_head(h[mx.arange(len(seqs)), self._last]).astype(mx.float32)
+        except _StopForward:
+            pass
+        captured = {k: v.astype(mx.float32) for k, v in self._captured.items()}
+        mx.eval(list(captured.values()) + ([out] if out is not None else []))
+        return ({k: np.array(v) for k, v in captured.items()},
+                np.array(out) if out is not None else None)
+
+
+def _repeatable(prefix: _Prefix) -> bool:
+    """A plain KV cache can be copied once per row of a batch; a recurrent state or a rotating window
+    goes through `forward`, one query at a time."""
+    return prefix.cache is not None and all(type(c) is KVCache for c in prefix.cache)