inventio 0.4.0__py3-none-any.whl

inventio/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.4.0"

inventio/_systemone/__init__.py

@@ -0,0 +1,128 @@
+"""Kev's serving path, vendored so a System One checkpoint needs one install and no second repository.
+
+Source: https://github.com/jaredpalmer/kev at f1535963cea021439370c23127bc970b6788e730 (2026-09-26),
+Apache-2.0, (c) the Kev authors. Copied: `model.py`, `checkpoint.py`, `api.py`, `device.py` — the four
+files `kev.serve` turns a checkpoint into answered questions with (Kev's own Hugging Face Space vendors
+the first three next to each other, so the layout is Kev's, not ours). Everything a serving pass
+touches is their code unchanged: the packed block-causal encode, the pointer head and its temperature,
+the readouts, the request/answer schema.
+
+Changed here, and nothing else:
+
+- `api.py`: the option cap is 512, Kev ships 255 (a code pool's median line count is 407; see
+  `inventio/systemone.py` for what 255 cost on SWE-bench Lite).
+- `model.py`: the CUDA-graph serving batch path (`kev.cuda_graphs`, `EAGER_STATES`) is gone;
+  `probs_batch` runs Kev's eager path, the one every measured number in this repository used.
+- `checkpoint.py`: the trainer-only and other-backend parts are gone — `mlx_available`,
+  `fused_available`, `backend`, `_load_mlx`, `warm_start`, `_load_backbone_into`,
+  `_load_adapter_into`, `weights_sha256`, `release_date`, `hybrid_base`, `LoadOptions.from_env` and
+  the `backend`/`cuda_graphs`/`fused` fields. `_load_torch` is the torch path alone. `resolve_run` reads a
+  Hub checkpoint from the cache when it is there instead of asking the Hub for the current revision on every
+  load (inventio's release check is `inventio update`'s, once a day). `Checkpoint.load` reads a full-weight
+  checkpoint's own tokenizer when it carries one, so the published model needs neither the base repo nor the network.
+- `device.py`: as published.
+- This file: `load()` (the device and dtype defaults `kev.serve` serves with, the Windows attention
+  fallback, and the dependency check that names the missing extra instead of raising an ImportError
+  from inside a forward pass).
+
+`load()` is the only entry point inventio uses; a checkpoint answers to any model name, so what
+identifies weights is the run directory or Hub id it was loaded from.
+"""
+
+import os
+import re
+
+__all__ = ["Checkpoint", "LoadOptions", "REQUIRES", "is_hub_id", "load", "missing", "patch_attention"]
+
+# What a checkpoint needs, and what an environment without it has to be told. `transformers>=5.17` is
+# the version the hybrid (Qwen3.5) cache classes came from; an older one imports and then fails inside
+# the first forward pass, which is a worse error than this one.
+REQUIRES = ("torch>=2.6", "transformers>=5.17,<6", "peft>=0.21", "pydantic>=2.9")
+
+
+def missing() -> str | None:
+    """What this environment still has to install to load a checkpoint, or None when it can.
+
+    Reads installed distributions, never imports them: `default_ranker()` asks this while the command line
+    is being built, and importing torch there would make every `inventio` command pay for a model runtime
+    it may not use. A package that answers here and then fails to import is caught at load time, where the
+    same message comes back.
+    """
+    import importlib.util
+    from importlib.metadata import version
+
+    for name, need in (("torch", "2.6"), ("transformers", "5.17"), ("peft", "0.21"), ("pydantic", "2.9")):
+        if importlib.util.find_spec(name) is None:
+            return f"{name} is not installed — pip install 'inventio[dispositio]'"
+    try:
+        if tuple(int(x) for x in re.findall(r"\d+", version("transformers"))[:2]) < (5, 17):
+            return (f"transformers {version('transformers')} is older than 5.17 (the hybrid cache classes "
+                    f"the encoder uses) — pip install -U 'inventio[dispositio]'")
+    except Exception:   # a distribution without version metadata: let the import speak
+        pass
+    return None
+
+
+def patch_attention() -> None:
+    """Windows: torch's SDPA GQA path falls back to the math kernel for the hybrid (Qwen3.5) bases, which
+    holds the whole L x L score matrix — out of memory past ~8k tokens, twice the time at 3k. Repeating
+    the kv heads instead lets SDPA take its memory-efficient kernel; measured on this machine as the
+    same numbers, a different kernel (benchmarks/kev_win.py, where the patch was found)."""
+    if os.name == "nt":
+        import transformers.integrations.sdpa_attention as sdpa
+
+        sdpa.use_gqa_in_sdpa = lambda *a, **k: False
+
+
+def load(run, device=None, *, dtype=None, merge: bool = True, attn: str | None = None,
+         lora_scale: float = 1.0, temperature: float | None = None):
+    """A checkpoint as `kev.serve` loads it -> (tokenizer, model).
+
+    What `kev.serve` serves with: bf16 on an accelerator (half the memory; probabilities within ~0.01 of
+    the exact path, the same argmax), fp32 on CPU; a checkpoint trained on a bf16 backbone loads in bf16
+    wherever it runs, as it was trained. INVENTIO_DTYPE=fp32 asks for the exact path — the one
+    every number in this repository was reported from — which costs about twice the wall time.
+    INVENTIO_DEVICE=cpu|cuda|mps leaves the accelerator to something else.
+    """
+    from .device import default_device
+
+    dev = str(device or os.environ.get("INVENTIO_DEVICE") or default_device())
+    if dev == "cpu":
+        # The hybrid (Qwen3.5) DeltaNet layers take a compiled kernel when one can be imported: from the
+        # Hugging Face kernel hub, or from flash-linear-attention (`fla`) and `causal_conv1d` if those are
+        # installed. Each is a GPU kernel, and on CPU it raises inside a Triton launch ("cannot be accessed
+        # from Triton (cpu tensor?)", measured with fla in the environment). transformers picks the
+        # implementation when its model module is imported, so the pure-torch path has to be asked for
+        # before that: the hub flag off, and the two packages made unimportable in this process.
+        import sys
+
+        os.environ.setdefault("USE_HUB_KERNELS", "NO")
+        for pkg in ("fla", "causal_conv1d"):
+            sys.modules.setdefault(pkg, None)   # `import fla` now raises ImportError: transformers falls back
+        if "transformers.models.qwen3_5.modeling_qwen3_5" in sys.modules and sys.modules.get("fla"):
+            raise RuntimeError("this process already loaded the GPU kernels for the model's DeltaNet layers; "
+                               "run on CPU in a fresh process (INVENTIO_DEVICE=cpu)")
+    miss = missing()
+    if miss:
+        raise RuntimeError(miss)
+    from .checkpoint import Checkpoint, LoadOptions
+
+    import torch
+
+    patch_attention()
+    device = dev
+    if dtype is None and device != "cpu" and os.environ.get("INVENTIO_DTYPE", "").lower() != "fp32":
+        dtype = torch.bfloat16
+    opts = LoadOptions(dtype=dtype, merge=merge, attn=attn, lora_scale=lora_scale, temperature=temperature)
+    return Checkpoint(run).load(device, opts)
+
+
+from .ids import is_hub_id   # noqa: E402  (torch-free: `inventio model --use` validates a name without the extra)
+
+
+def __getattr__(name):   # PEP 562: Checkpoint/LoadOptions/resolve_run pull torch in, so they load on first use
+    if name in ("Checkpoint", "LoadOptions", "resolve_run"):
+        from . import checkpoint
+
+        return getattr(checkpoint, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

inventio/_systemone/api.py

@@ -0,0 +1,165 @@
+"""TypeSafe-compatible request/response shapes (POST /v1/systemone) mapped onto the single pointer primitive.
+
+Noul   -> 2 options [false, true];            answer = p(true)
+Choice -> options 'name' or 'name: desc';      answer = argmax, probabilities by name, confidence
+Score  -> options = ordered level descriptions; answer = expected level, legend, probabilities by index
+"""
+import json
+import re
+from datetime import datetime
+from typing import Any, Literal, Union
+from pydantic import BaseModel, Field, model_validator
+
+JSONContent = Union[str, dict, list, int, float, bool, None]
+MAX_OPTIONS = 512   # Kev ships 255; a code pool's median line count is 407 (inventio/systemone.py)
+
+
+class Noul(BaseModel):
+    type: Literal["noul"]
+    instructions: JSONContent = None
+    criteria: dict[str, JSONContent] | None = None
+
+
+class Choice(BaseModel):
+    type: Literal["choice"]
+    instructions: JSONContent = None
+    criteria: dict[str, JSONContent]
+
+    @model_validator(mode="after")
+    def _check(self):
+        if not 1 <= len(self.criteria) <= MAX_OPTIONS: raise ValueError(f"criteria must have 1..{MAX_OPTIONS} options")
+        return self
+
+
+class Score(BaseModel):
+    type: Literal["score"]
+    instructions: JSONContent = None
+    criteria: list[JSONContent] = Field(min_length=1, max_length=MAX_OPTIONS)
+
+
+Question = Union[Noul, Choice, Score]
+
+
+class SystemOneRequest(BaseModel):
+    state: JSONContent
+    model: str = "kev-latest"
+    questions: dict[str, Question] = Field(min_length=1)
+
+
+def render(v: JSONContent, indent: int = 0) -> str:
+    """Flatten str | object | array into text the model sees. Field names are kept as labels."""
+    pad = "  " * indent
+    if v is None: return ""
+    if isinstance(v, (str, int, float, bool)): return str(v)
+    if isinstance(v, list): return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
+    return "\n".join(f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}" for k, x in v.items())
+
+
+def option_text(name: str, desc: JSONContent) -> str:
+    return name if desc is None or desc == "" else f"{name}: {render(desc)}"
+
+
+MONTHS = "January|February|March|April|May|June|July|August|September|October|November|December"
+_DATE = re.compile(rf"\b(?:{MONTHS}) \d{{1,2}}, \d{{4}}\b|\b\d{{4}}-\d{{2}}-\d{{2}}\b")
+
+
+def date_facts(text: str) -> str:
+    """Deterministic date arithmetic for the model: every pair of absolute dates found in `text`, as one sentence each
+    ("August 3, 2026 is 12 days after July 22, 2026."). The model cannot subtract dates reliably (issue #8); it can use a
+    stated day count. Returns "" when fewer than two dates are found. Dates are listed in order of first appearance."""
+    found = []
+    for m in _DATE.finditer(text):
+        raw = m.group(0)
+        try: d = datetime.strptime(raw, "%B %d, %Y") if "," in raw else datetime.strptime(raw, "%Y-%m-%d")
+        except ValueError: continue
+        if raw not in [r for r, _ in found]: found.append((raw, d))
+    facts = []
+    for i in range(len(found)):
+        for j in range(i + 1, len(found)):
+            n = (found[j][1] - found[i][1]).days
+            facts.append(f"{found[j][0]} is {abs(n)} day{'s' if abs(n) != 1 else ''} {'after' if n > 0 else 'before'} {found[i][0]}." if n else f"{found[j][0]} is the same day as {found[i][0]}.")
+    return " ".join(facts)
+
+
+def with_date_facts(state):
+    """State with a `date_facts` field (object states) or an appended paragraph (string states) when two or more absolute
+    dates appear. Opt-in preprocessing (KEV_DATE_FACTS=1 in kev.serve, --date_facts in kev.benchmark)."""
+    facts = date_facts(render(state))
+    if not facts: return state
+    if isinstance(state, dict): return {**state, "date_facts": facts}
+    if isinstance(state, list): return state + [{"date_facts": facts}]
+    return f"{state}\n\ndate_facts: {facts}"
+
+
+def question_keys(qtype: str, criteria) -> list[str]:
+    """The keys a question's probabilities are reported under, in option order: the criteria names (choice),
+    ["false", "true"] (noul), the level indices as strings (score). Labels, targets and anchors use the same keys."""
+    if qtype == "choice": return list(criteria)
+    if qtype == "noul": return ["false", "true"]
+    return [str(i) for i in range(len(criteria))]
+
+
+def to_record(req: SystemOneRequest):
+    """-> internal record for encode(), plus per-question metadata ({"id", "type", "keys", "legend" for score}) to map
+    probabilities back."""
+    qs, meta = [], []
+    for qid, q in req.questions.items():
+        m = {"id": qid, "type": q.type, "keys": question_keys(q.type, q.criteria)}
+        if q.type == "noul":
+            c = q.criteria or {}
+            opts = [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
+        elif q.type == "choice":
+            opts = [option_text(k, v) for k, v in q.criteria.items()]
+        else:
+            opts = [render(x) for x in q.criteria]
+            m["legend"] = dict(zip(m["keys"], opts))
+        qs.append({"instr": render(q.instructions), "options": opts, "label": 0}); meta.append(m)
+    return {"state": render(req.state), "questions": qs}, meta
+
+
+# Both confidence formulas mirror TypeSafe's reference adapter, system-one-adapter 0.2.1
+# (src/system_one_adapter/_utils/confidence_metrics.py): p is normalised to sum 1 (all zeros -> uniform), one option -> 1.
+def _normalize(p: list[float]) -> list[float]:
+    t = sum(p)
+    return [1 / len(p)] * len(p) if t == 0 else [x / t for x in p]
+
+
+def choice_confidence(p: list[float]) -> float:
+    """(p_max - 1/K) / (1 - 1/K): 0 at uniform, 1 at certainty."""
+    K = len(p)
+    return 1.0 if K == 1 else (max(_normalize(p)) - 1 / K) / (1 - 1 / K)
+
+
+def score_confidence(p: list[float]) -> float:
+    """max(0, 1 - E|level - mode| / D), D = mean absolute deviation of a uniform distribution over the L levels around its
+    mean (L-1)/2; mode = first most likely level. 1 when all mass is on one level, 0 at uniform or anything as spread."""
+    L = len(p)
+    if L == 1: return 1.0
+    p = _normalize(p); mode = max(range(L), key=p.__getitem__)
+    D = sum(abs(i - (L - 1) / 2) for i in range(L)) / L
+    return max(0.0, 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / D)
+
+
+def round_prob(x: float) -> float:
+    """Serialization precision for probabilities and derived scalars. 4 decimals keeps the sum of a rounded distribution
+    within TypeSafe's tolerance (|sum - 1| < 0.02) at the 512-option maximum: 512 * 0.00005 < 0.02."""
+    return round(float(x), 4)
+
+
+def to_answers(probs: list[list[float]], meta: list[dict]) -> dict[str, Any]:
+    out = {}
+    for p, m in zip(probs, meta):
+        if m["type"] == "noul":
+            out[m["id"]] = {"type": "noul", "noul": round_prob(p[1])}
+        elif m["type"] == "choice":
+            dist = {k: round_prob(v) for k, v in zip(m["keys"], p)}
+            out[m["id"]] = {"type": "choice", "choice": m["keys"][max(range(len(p)), key=lambda i: p[i])], "confidence": round_prob(choice_confidence(p)), "probabilities": dist}
+        else:
+            score = sum(i * pi for i, pi in enumerate(p))
+            out[m["id"]] = {"type": "score", "score": round_prob(score), "legend": m["legend"], "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)}, "confidence": round_prob(score_confidence(p))}
+    return out
+
+
+def output_tokens(tok, answers: dict) -> int:
+    """Billing-style figure: tokens of the serialised answers. Not a measure of generation (there is none)."""
+    return len(tok(json.dumps(answers), add_special_tokens=False).input_ids)

inventio/_systemone/checkpoint.py

@@ -0,0 +1,192 @@
+"""Trained checkpoints: a run directory or a Hub repo holding a LoRA adapter (or, for a full-weight run, the whole bf16
+backbone), `head.pt` and the tokenizer.
+
+Loader rule: `adapter_config.json` present -> a LoRA adapter on `meta.base` at `meta.base_revision`; no adapter and
+`config.json` + `model*.safetensors` (save_pretrained of the backbone, `meta.weights == "full"`) -> the backbone is loaded from
+the checkpoint directory itself, nothing is merged, in the dtype head.pt's `weights_dtype` names (it must match the `dtype`
+save_pretrained wrote to config.json). The tokenizer always comes from the base (both layouts carry a copy).
+
+This is the one place that knows the layout of `head.pt` and how a checkpoint becomes a `DecisionModel`:
+`kev.serve`, `kev.benchmark`, `kev.train --init_from`, `kev.publish`, the scripts and the Hugging Face Space all go
+through it. The Space vendors this file next to `model.py` and `api.py` (scripts/publish_space.sh), so it must not
+import the data or suite modules at import time.
+
+    ck = Checkpoint("jaredpalmer/kev-4b")          # or a local run directory; `@tag` pins a Hub revision
+    tok, model = ck.load("mps", LoadOptions.from_env())
+    ck.meta.temperature                             # the calibration the checkpoint carries
+"""
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+
+from .model import DecisionModel, load_tokenizer
+
+from .ids import is_hub_id   # vendored here: no torch needed to answer a question about a name
+
+
+def resolve_run(run):
+    """Local run directory as given, or a Hub repo id like jaredpalmer/kev-4b, optionally pinned to a revision or tag
+    with `@` (jaredpalmer/kev-4b@qwen3), downloaded to the HF cache. Returns a str path."""
+    if os.path.isdir(run):
+        return str(run)
+    from huggingface_hub import snapshot_download
+    repo, _, revision = str(run).partition("@")
+    pats = ["*.json", "*.safetensors", "*.pt", "*.txt", "*.jinja"]
+    try:   # the cache first: a load must not ask the network which revision is current (inventio: `update` does)
+        return snapshot_download(repo, revision=revision or None, allow_patterns=pats, local_files_only=True)
+    except Exception:   # not downloaded yet: the first load fetches it
+        return snapshot_download(repo, revision=revision or None, allow_patterns=pats)
+
+
+@dataclass
+class Meta:
+    """Contents of `head.pt`. Every reader gets the same defaults for fields older checkpoints did not write.
+    `extra` keeps the rest of the file (training args, suite hash, init provenance, temperature fit) so a
+    read-modify-write round trip loses nothing."""
+    base: str
+    head: dict | None = None
+    base_revision: str | None = None
+    lora: int = 0
+    head_dim: int = 256
+    option_isolation: bool = False
+    special_embeddings: bool = False
+    weights_dtype: str = "fp32"
+    temperature: float = 1.0
+    holdout: list = field(default_factory=list)
+    weights: str = "lora"          # "lora": an adapter on the base; "full": the whole backbone is in the checkpoint (kev.train --full_ft)
+    extra: dict = field(default_factory=dict)
+
+    KNOWN = ("base", "head", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights_dtype", "temperature", "holdout", "weights")
+
+    @classmethod
+    def from_dict(cls, d):
+        return cls(**{k: d[k] for k in cls.KNOWN if k in d}, extra={k: v for k, v in d.items() if k not in cls.KNOWN})
+
+    def to_dict(self):
+        return {**self.extra, **{k: getattr(self, k) for k in self.KNOWN}}   # known fields win over a stray key in extra
+
+
+def read_meta(run):
+    return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu"))
+
+
+def write_meta(run, meta):
+    torch.save(meta.to_dict(), f"{run}/head.pt")
+
+
+@dataclass(frozen=True)
+class LoadOptions:
+    """How a checkpoint is turned into a model. Defaults are the exact path every reported number uses.
+
+    dtype        None = fp32, the exact path every reported number uses (bf16 when the checkpoint was trained with a bf16
+                 backbone). Serving defaults to bf16 on CUDA and MPS instead: half the memory, 2-4.5x lower latency on an
+                 L4 (Kev-4B: 209 -> 118 ms at 101 tokens, 850 -> 189 ms at 330 tokens), probabilities within ~0.01 and
+                 the same argmax on the checks run so far. The exact path is INVENTIO_DTYPE=fp32.
+    merge        fold the LoRA into the base weights: the delta is computed from the fp32 adapter and added in fp32 with
+                 one rounding to the load dtype, so a bf16 model holds exactly round(W + delta), the same bits as merging
+                 an fp32 copy and casting, without the fp32 copy (Kev-9B needed 36 GB of GPU memory to load for that).
+                 Identical because the Qwen bases are stored in bf16; a base stored in fp32 would be rounded twice.
+                 Exact in fp32; in bf16 it is faster (~15%) and closer to the fp32 numbers than the unmerged adapter
+                 (kev-4b, 24 dev records: max |dp| 0.017 vs 0.029, 0 vs 1 argmax flips). Ignored for adapters that carry
+                 trained token embeddings. A checkpoint trained on a bf16 backbone keeps its adapter unmerged, as it was
+                 trained (the fused serving path that folded it is not vendored).
+    attn         attention backend; None = the model default (SDPA on CUDA, eager elsewhere). "sdpa" on MPS measured
+                 parity with eager and is a few percent faster.
+    lora_scale   WiSE-FT-style interpolation between base (0) and fine-tuned weights (1), at inference.
+    temperature  None = the temperature the checkpoint carries (fitted by scripts/calibrate_checkpoint.py); 1.0 = raw logits.
+    """
+    dtype: torch.dtype | None = None
+    merge: bool = True
+    attn: str | None = None
+    lora_scale: float = 1.0
+    temperature: float | None = None
+
+
+
+class Checkpoint:
+    def __init__(self, run):
+        self.requested = str(run)                    # what the caller asked for (a Hub id stays a Hub id in labels)
+        self.path = resolve_run(run)
+        self.meta = read_meta(self.path)
+
+    def file(self, name):
+        return Path(self.path) / name
+
+    def adapter_config(self):
+        return json.loads(self.file("adapter_config.json").read_text(encoding="utf-8"))
+
+    @property
+    def full(self):
+        """The loader rule (module docstring): True for a full-weight checkpoint, False for a LoRA adapter; head.pt's
+        `weights` must agree with the files."""
+        found = "lora" if self.file("adapter_config.json").exists() else "full" if self.file("config.json").exists() and self.shards() else None
+        if found != self.meta.weights:
+            raise ValueError(f"{self.path}: head.pt says weights={self.meta.weights!r} but the directory holds "
+                             f"{ {'lora': 'an adapter', 'full': 'backbone weights'}.get(found, 'neither an adapter nor backbone weights') }")
+        return found == "full"
+
+    def shards(self):
+        """The backbone's safetensors files of a full-weight checkpoint (model.safetensors or model-*-of-*.safetensors)."""
+        return sorted(Path(self.path).glob("model*.safetensors"))
+
+    def load(self, device, opts=LoadOptions()):
+        """-> (tokenizer, model) in eval mode with the LoRA applied (or the full backbone loaded) and the pointer head
+        loaded. The model is a DecisionModel, ready for encode()/probs()."""
+        meta = self.meta
+        # a checkpoint that carries its tokenizer is read from itself: a published full-weight model then loads with
+        # no second repository and no network (inventio's offline promise); otherwise the base's, as Kev does
+        own = self.file("tokenizer.json").exists() and self.full
+        tok = load_tokenizer(self.path) if own else load_tokenizer(meta.base, revision=meta.base_revision)
+        m = self._load_torch(tok, device, opts)
+        m.head.load_state_dict(meta.head); m.eval()
+        m.head.temperature = meta.temperature if opts.temperature is None else opts.temperature
+        return tok, m
+
+    def _load_torch(self, tok, device, opts):
+        return self._full_torch(tok, device, opts)[0] if self.full else self._adapted_torch(tok, device, opts)[0]
+
+    SAVED_DTYPES = {"bf16": "bfloat16", "fp32": "float32"}   # head.pt weights_dtype -> the dtype save_pretrained writes to config.json
+
+    def _full_torch(self, tok, device, opts):
+        """-> (model, True). Full weights load in the dtype head.pt's `weights_dtype` names (bf16 for every kev.train
+        --full_ft run: the dtype they were trained in), which must be the dtype save_pretrained recorded in config.json;
+        otherwise a mislabelled export would be silently cast (fp32 weights rounded to bf16, or bf16 upcast to twice the
+        memory). An explicit dtype still casts on purpose (fp32: the same values computed in fp32). Nothing to merge."""
+        if opts.lora_scale != 1: raise ValueError("lora_scale interpolates an adapter; a full-weight checkpoint has none")
+        meta = self.meta
+        cfg = json.loads(self.file("config.json").read_text(encoding="utf-8"))
+        expected, saved = self.SAVED_DTYPES.get(meta.weights_dtype), cfg.get("dtype") or cfg.get("torch_dtype")
+        if expected is None or saved not in (None, expected):
+            raise ValueError(f"{self.path}: config.json records the weights as {saved} but head.pt says weights_dtype={meta.weights_dtype!r}")
+        return DecisionModel(meta.base, tok, device, head_dim=meta.head_dim, option_isolation=meta.option_isolation,
+                             dtype=opts.dtype or getattr(torch, expected), attn=opts.attn, weights=self.path), True
+
+    def _adapted_torch(self, tok, device, opts):
+        """-> (model, whether the adapter was merged): the base with this checkpoint's LoRA."""
+        from peft import PeftModel
+        meta = self.meta
+        dtype, merge = opts.dtype or torch.float32, opts.merge
+        if meta.weights_dtype == "bf16":
+            # trained with a bf16 backbone: load it the same way. The exact path keeps the fp32 adapter unmerged.
+            dtype, merge = torch.bfloat16, merge
+        merge = merge and not self.adapter_config().get("trainable_token_indices")   # token-trained adapters stay unmerged
+        m = DecisionModel(meta.base, tok, device, lora=None, revision=meta.base_revision, head_dim=meta.head_dim,
+                          option_isolation=meta.option_isolation, dtype=dtype, attn=opts.attn)
+        m.lm = PeftModel.from_pretrained(m.lm, self.path, torch_device=str(device)).to(device)   # trainable token embeddings, if any, live in the adapter
+        if opts.lora_scale != 1:
+            for module in m.lm.modules():
+                if isinstance(getattr(module, "scaling", None), dict):
+                    for k in module.scaling: module.scaling[k] *= opts.lora_scale
+            m.lora_scale = opts.lora_scale
+        if merge: m.lm = m.lm.merge_and_unload()     # W += delta: fp32 math, one rounding (see LoadOptions.merge)
+        if dtype != torch.float32: m.lm = m.lm.to(dtype)
+        return m, merge
+
+    COMPAT_FIELDS = ("base", "base_revision", "lora", "head_dim", "option_isolation", "special_embeddings", "weights")
+
+def load(run, device, opts=LoadOptions()):
+    """Convenience: Checkpoint(run).load(device, opts)."""
+    return Checkpoint(run).load(device, opts)

inventio/_systemone/device.py

@@ -0,0 +1,30 @@
+"""The accelerator this process uses: cuda, then mps, then cpu."""
+import torch
+
+
+def default_device():
+    return "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+def sync(device):
+    """Wait for queued kernels, so wall-clock timings around a forward pass are real."""
+    if device == "mps": torch.mps.synchronize()
+    elif device == "cuda": torch.cuda.synchronize()
+
+
+def empty_cache(device):
+    if device == "mps": torch.mps.empty_cache()
+    elif device == "cuda": torch.cuda.empty_cache()
+
+
+def out_of_memory(e):
+    """Whether a torch allocator ran out: CUDA raises torch.OutOfMemoryError, MPS a plain RuntimeError with this message.
+    The MLX backend's Metal errors are neither."""
+    return isinstance(e, torch.OutOfMemoryError) or isinstance(e, RuntimeError) and str(e).startswith("MPS backend out of memory")
+
+
+def allocated_bytes(device):
+    """Bytes currently allocated on the device (MPS) or the peak since the process started (CUDA); 0 on CPU."""
+    if device == "mps": return torch.mps.current_allocated_memory()
+    if device == "cuda": return torch.cuda.max_memory_allocated()
+    return 0

inventio/_systemone/ids.py

@@ -0,0 +1,15 @@
+"""What a checkpoint can be named by: a directory, or a Hub repo id (optionally pinned `@revision`).
+
+Kev's `checkpoint.py` keeps the same rule; it lives here too because `inventio model --use` has to
+validate a name on a machine that has not installed the extra, and importing `checkpoint` would import
+torch and transformers to answer a question about a string.
+"""
+
+import os
+import re
+
+HUB_ID = re.compile(r"[\w.-]+/[\w.-]+(@[\w.-]+)?")
+
+
+def is_hub_id(run) -> bool:
+    return not os.path.isdir(run) and HUB_ID.fullmatch(str(run)) is not None