sysone 0.2.0__py3-none-any.whl

sysone/__init__.py

@@ -0,0 +1,22 @@
+"""Fine-tune and calibrate system-one decision models with a scikit-learn style fit / predict interface."""
+
+from .architectures import Architecture, register_architecture
+from .base import NotFittedError, SystemOneModel
+from .cache import cache_info, clear_cache
+from .evaluation import baselines, compare
+from .models import DecoderModel, EncoderModel
+from .version import __version__
+
+__all__ = [
+    "Architecture",
+    "DecoderModel",
+    "EncoderModel",
+    "NotFittedError",
+    "SystemOneModel",
+    "__version__",
+    "baselines",
+    "cache_info",
+    "clear_cache",
+    "compare",
+    "register_architecture",
+]

sysone/_util.py

@@ -0,0 +1,102 @@
+"""Small helpers shared across modules."""
+
+from __future__ import annotations
+
+import contextlib
+import math
+from numbers import Integral, Real
+
+import torch
+
+
+def resolve_count(value, total: int, name: str) -> int:
+    """Resolve a layer count given as an int, a fraction in (0, 1], or "all".
+
+    Follows scikit-learn's convention: an int is a count and a float is a fraction, so `1` is one
+    layer and `1.0` is all of them. A fraction is rounded to the nearest whole layer, and a fraction
+    above 0 gives at least one.
+    """
+    if value is None or value == "all":
+        return total
+    if isinstance(value, bool):
+        raise ValueError("%s must be an int, a fraction in (0, 1] or 'all', got %r" % (name, value))
+    if isinstance(value, Integral):
+        if not 0 <= value <= total:
+            raise ValueError("%s=%d is outside 0..%d (the model has %d layers here)" % (name, value, total, total))
+        return int(value)
+    if isinstance(value, Real):
+        if not 0.0 <= value <= 1.0 or math.isnan(value):
+            raise ValueError("%s as a fraction must be in [0, 1], got %r" % (name, value))
+        if value == 0:
+            return 0
+        return max(1, min(total, int(round(value * total))))
+    raise ValueError("%s must be an int, a fraction in (0, 1] or 'all', got %r" % (name, value))
+
+
+def resolve_device(device) -> torch.device:
+    if device in (None, "auto"):
+        if torch.cuda.is_available():
+            return torch.device("cuda")
+        if torch.backends.mps.is_available():
+            return torch.device("mps")
+        return torch.device("cpu")
+    return torch.device(device)
+
+
+_PRECISIONS = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def resolve_dtype(precision: str, device: torch.device) -> torch.dtype:
+    """The dtype for frozen backbone weights. "auto" is bf16 where the device supports it."""
+    if precision == "auto":
+        if device.type == "cuda" and torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+        if device.type == "mps":
+            return torch.bfloat16
+        return torch.float32
+    if precision not in _PRECISIONS:
+        raise ValueError("precision must be one of auto, %s, got %r" % (", ".join(_PRECISIONS), precision))
+    return _PRECISIONS[precision]
+
+
+@contextlib.contextmanager
+def quiet_transformers():
+    """Silence transformers' load reports (e.g. weights of layers dropped by num_layers)."""
+    from transformers.utils import logging
+
+    previous = logging.get_verbosity()
+    bars = logging.is_progress_bar_enabled()
+    logging.set_verbosity_error()
+    logging.disable_progress_bar()
+    try:
+        yield
+    finally:
+        logging.set_verbosity(previous)
+        if bars:
+            logging.enable_progress_bar()
+
+
+def is_out_of_memory(exc: BaseException) -> bool:
+    return isinstance(exc, torch.OutOfMemoryError) or "out of memory" in str(exc).lower()
+
+
+def empty_device_cache() -> None:
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    if torch.backends.mps.is_available():
+        torch.mps.empty_cache()
+
+
+def autocast(device: torch.device, dtype: torch.dtype | None):
+    """Autocast to `dtype` on `device`, or nothing when `dtype` is None."""
+    if dtype is None:
+        return contextlib.nullcontext()
+    return torch.autocast(device_type=device.type, dtype=dtype)
+
+
+def format_bytes(n: float) -> str:
+    for unit in ("B", "KB", "MB", "GB"):
+        if abs(n) < 1024 or unit == "GB":
+            return ("%d %s" if unit == "B" else "%.1f %s") % (n, unit)
+        n /= 1024
+    return "%.1f GB" % n

sysone/answers.py

@@ -0,0 +1,57 @@
+"""Turn option probabilities into answers in the Jev `/v1/systemone` shape."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from .questions import Question
+
+_DIGITS = 4
+
+
+def choice_confidence(p: Sequence[float]) -> float:
+    """`(max p - 1/n) / (1 - 1/n)`: 0 for a uniform distribution, 1 for a certain one."""
+    n = len(p)
+    if n < 2:
+        return 1.0
+    return max(0.0, (max(p) - 1.0 / n) / (1.0 - 1.0 / n))
+
+
+def score_confidence(p: Sequence[float]) -> float:
+    """One minus the spread around the mode, normalised by the spread of a uniform distribution."""
+    n = len(p)
+    if n < 2:
+        return 1.0
+    mode = max(range(n), key=p.__getitem__)
+    spread = sum(v * abs(i - mode) for i, v in enumerate(p))
+    center = (n - 1) / 2
+    uniform_spread = sum(abs(i - center) for i in range(n)) / n
+    return max(0.0, 1.0 - spread / uniform_spread)
+
+
+def answer(question: Question, p: Sequence[float]) -> dict:
+    """One answer dict for `question` from its option probabilities (in option order)."""
+    p = [float(v) for v in p]
+    top = max(p)
+    if question.type == "choice":
+        return {
+            "type": "choice",
+            "choice": question.keys[p.index(top)],
+            "probabilities": {str(key): round(v, _DIGITS) for key, v in zip(question.keys, p)},
+            "confidence": round(choice_confidence(p), _DIGITS),
+            "answer_confidence": round(top, _DIGITS),
+        }
+    if question.type == "score":
+        return {
+            "type": "score",
+            "score": round(sum(i * v for i, v in enumerate(p)), _DIGITS),
+            "legend": {str(i): text for i, text in enumerate(question.texts)},
+            "probabilities": {str(i): round(v, _DIGITS) for i, v in enumerate(p)},
+            "confidence": round(score_confidence(p), _DIGITS),
+            "answer_confidence": round(top, _DIGITS),
+        }
+    return {
+        "type": "noul",
+        "noul": round(p[1], _DIGITS),
+        "answer_confidence": round(top, _DIGITS),
+    }

sysone/architectures/__init__.py

@@ -0,0 +1,12 @@
+from .base import Architecture
+from .modernbert import ModernBertArchitecture
+from .qwen import QwenArchitecture
+from .registry import architecture_for, register_architecture
+
+__all__ = [
+    "Architecture",
+    "ModernBertArchitecture",
+    "QwenArchitecture",
+    "architecture_for",
+    "register_architecture",
+]

sysone/architectures/base.py

@@ -0,0 +1,144 @@
+"""The Architecture interface: everything model-specific sits behind it."""
+
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+
+import torch
+from torch import nn
+
+from .._util import quiet_transformers, resolve_count
+from ..layout import Pieces, attention_mask
+
+
+def encode_text(tokenizer, text: str) -> list[int]:
+    """Token ids for one piece of text, without special tokens. Pieces are joined afterwards, so the
+    exact token span of every option is known."""
+    return tokenizer(text, add_special_tokens=False)["input_ids"]
+
+
+class Architecture(ABC):
+    """How to use one model family.
+
+    Subclasses say how pieces are ordered and which positions are read (`pieces`). Loading, layer
+    access, masks and the forward pass have defaults that suit Hugging Face models whose base
+    model has a `layers` list and accepts masks as a dict keyed by layer type.
+    """
+
+    #: `AutoConfig.model_type` values this architecture handles.
+    model_types: tuple[str, ...] = ()
+    #: True for left-to-right models, whose tokens only see what came before them.
+    causal: bool = False
+    #: Module paths (regex fragments under a layer) for the named adapter target groups.
+    lora_targets: dict[str, list[str]] = {}
+    #: `backbone_lr="auto"`: the learning rate for fully trained layers ("full") and for adapters
+    #: ("adapter"). Subclasses set values measured on their model family.
+    learning_rates: dict[str, float] = {"full": 2.5e-5, "adapter": 1e-4}
+
+    # -- loading -----------------------------------------------------------------------------
+
+    def load(self, name: str, num_layers, dtype: torch.dtype, revision: str | None = None,
+             quantization=None, device: torch.device | None = None) -> nn.Module:
+        """Load the base model with only the first `num_layers` layers built.
+
+        `quantization` is a transformers `BitsAndBytesConfig`; a quantized model is placed on
+        `device` while loading, since it can't be moved afterwards.
+
+        Always uses `sdpa` attention: flash attention supports only standard causal or padding
+        masks and would silently ignore the layout's masks.
+        """
+        from transformers import AutoConfig, AutoModel
+
+        config = AutoConfig.from_pretrained(name, revision=revision)
+        total = config.num_hidden_layers
+        keep = resolve_count(num_layers, total, "num_layers")
+        if keep < 1:
+            raise ValueError("num_layers must keep at least one layer")
+        config.num_hidden_layers = keep
+        if getattr(config, "layer_types", None):
+            config.layer_types = list(config.layer_types[:keep])
+        extra = {}
+        if quantization is not None:
+            extra = {"quantization_config": quantization, "device_map": {"": device}}
+        with quiet_transformers():
+            model = AutoModel.from_pretrained(name, config=config, dtype=dtype, attn_implementation="sdpa",
+                                              revision=revision, **extra)
+        model.sysone_total_layers = total
+        return model
+
+    def layers(self, model: nn.Module) -> nn.ModuleList:
+        return model.layers
+
+    def embeddings(self, model: nn.Module) -> nn.Module:
+        """The token embeddings, trained with `train_embeddings=True`."""
+        return model.get_input_embeddings()
+
+    def adapter_modules(self, model: nn.Module, targets, layer_indices) -> list[str]:
+        """Full names of the linear modules in `layer_indices` that get adapters.
+
+        `targets` is a named group (`"attention"`, `"mlp"`, `"all"`) translated through
+        `lora_targets`, or a list of module names, matched by path within a layer (`"q_proj"` or
+        `"attn.Wo"`) and checked against the loaded model.
+        """
+        groups = {**self.lora_targets, "all": [p for g in ("attention", "mlp") for p in self.lora_targets.get(g, [])]}
+        if isinstance(targets, str):
+            if targets not in groups or not groups[targets]:
+                raise ValueError("lora_targets=%r is not supported by %s; use one of %s or a list of module names"
+                                 % (targets, self.name(), ", ".join(repr(g) for g in groups if groups[g])))
+            patterns = [re.compile(r"(^|\.)(%s)$" % p) for p in groups[targets]]
+        else:
+            patterns = [re.compile(r"(^|\.)%s$" % re.escape(t)) for t in targets]
+        layers = self.layers(model)
+        prefix = next(n for n, m in model.named_modules() if m is layers)
+        names, matched = [], set()
+        for i in layer_indices:
+            for sub, module in layers[i].named_modules():
+                if not isinstance(module, nn.Linear):
+                    continue
+                hits = [j for j, p in enumerate(patterns) if p.search(sub)]
+                if hits:
+                    names.append("%s.%d.%s" % (prefix, i, sub))
+                    matched.update(hits)
+        if not isinstance(targets, str) and len(matched) < len(patterns):
+            missing = [t for j, t in enumerate(targets) if j not in matched]
+            available = sorted(sub for sub, m in layers[0].named_modules() if isinstance(m, nn.Linear))
+            raise ValueError("lora_targets %s matched no module of %s; its layers have these linear modules: %s"
+                             % (missing, self.name(), ", ".join(available)))
+        return names
+
+    def max_positions(self, model: nn.Module) -> int:
+        return int(model.config.max_position_embeddings)
+
+    # -- sequence layout -----------------------------------------------------------------------
+
+    def check_tokenizer(self, tokenizer) -> None:
+        if tokenizer.pad_token_id is None:
+            raise ValueError("the tokenizer has no pad token")
+
+    @abstractmethod
+    def pieces(self, tokenizer, head: list[int], state: list[int], options: list[list[int]]) -> Pieces:
+        """Arrange already-tokenized pieces: `head` (task line and question), `state`, and each
+        option's own text. Adds the model's special tokens and delimiters."""
+
+    # -- forward -----------------------------------------------------------------------------
+
+    def attention(self, model: nn.Module, real, position_ids, option_ids) -> dict[str, torch.Tensor]:
+        """Masks keyed by layer type, the form the model's forward accepts."""
+        layer_types = set(getattr(model.config, "layer_types", None) or ["full_attention"])
+        masks = {}
+        for layer_type in layer_types:
+            window = model.config.sliding_window if layer_type == "sliding_attention" else None
+            masks[layer_type] = attention_mask(real, position_ids, option_ids, causal=self.causal,
+                                               sliding_window=window)
+        return masks
+
+    def hidden_states(self, model: nn.Module, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        masks = self.attention(model, batch["real"], batch["position_ids"], batch["option_ids"])
+        out = model(input_ids=batch["input_ids"], attention_mask=masks, position_ids=batch["position_ids"],
+                    use_cache=False) if self.causal else \
+            model(input_ids=batch["input_ids"], attention_mask=masks, position_ids=batch["position_ids"])
+        return out.last_hidden_state
+
+    def name(self) -> str:
+        return self.model_types[0] if self.model_types else type(self).__name__

sysone/architectures/modernbert.py

@@ -0,0 +1,43 @@
+"""ModernBERT: a bidirectional encoder."""
+
+from __future__ import annotations
+
+from ..layout import Pieces
+from .base import Architecture
+
+
+class ModernBertArchitecture(Architecture):
+    """`[CLS] head [SEP] [MASK] option … [MASK] option [SEP] state [SEP]`.
+
+    Each option is read at the `[MASK]` that opens it, which sees everything in a bidirectional
+    model. The summary is read at `[CLS]`.
+    """
+
+    model_types = ("modernbert",)
+    causal = False
+    lora_targets = {
+        # ModernBERT names both its attention output and its MLP output `Wo`, so targets are
+        # matched by path, never by bare module name.
+        "attention": [r"attn\.Wqkv", r"attn\.Wo"],
+        "mlp": [r"mlp\.Wi", r"mlp\.Wo"],
+    }
+    # Full tuning at 5e-5 reached 0.79 test accuracy on typed-decisions; adapters at 1e-4 barely
+    # moved in 4 epochs, so their default is ten times higher (not yet measured).
+    learning_rates = {"full": 5e-5, "adapter": 1e-3}
+
+    def check_tokenizer(self, tokenizer) -> None:
+        super().check_tokenizer(tokenizer)
+        missing = [name for name in ("cls_token_id", "sep_token_id", "mask_token_id")
+                   if getattr(tokenizer, name) is None]
+        if missing:
+            raise ValueError("the tokenizer has no %s" % ", ".join(missing))
+
+    def pieces(self, tokenizer, head, state, options) -> Pieces:
+        cls, sep, mask = tokenizer.cls_token_id, tokenizer.sep_token_id, tokenizer.mask_token_id
+        return Pieces(
+            prefix=[cls, *head, sep],
+            options=[[mask, *option] for option in options],
+            option_read=[0] * len(options),
+            suffix=[sep, *state, sep],
+            summary=("prefix", 0),
+        )

sysone/architectures/qwen.py

@@ -0,0 +1,48 @@
+"""Qwen3: a causal decoder."""
+
+from __future__ import annotations
+
+from ..layout import Pieces
+from .base import Architecture, encode_text
+
+OPEN, CLOSE, CUE = "<option>", "</option>", "The answer is:"
+
+
+class QwenArchitecture(Architecture):
+    """`head  state  <option>option</option> …  The answer is:`.
+
+    A causal token only sees what came before it, so:
+    - the question comes before the state, and every state token is read knowing what is asked;
+    - options come after the state, so they can see it;
+    - each option is read at the last token of its own text, just before `</option>`, the first
+      point at which it has seen the whole option. `<option>` / `</option>` are plain-text
+      delimiters split into ordinary pretrained tokens; no vector is read from them;
+    - the summary is the last token of the answer cue, which has seen everything.
+    """
+
+    model_types = ("qwen3",)
+    causal = True
+    lora_targets = {
+        "attention": [r"self_attn\.(q|k|v|o)_proj"],
+        "mlp": [r"mlp\.(gate|up|down)_proj"],
+    }
+    # Adapters at 1e-4 measured on typed-decisions; the full-tuning rate is not yet measured.
+    learning_rates = {"full": 2.5e-5, "adapter": 1e-4}
+
+    def pieces(self, tokenizer, head, state, options) -> Pieces:
+        gap, newline, open_, close, cue = _delimiters(tokenizer)
+        return Pieces(
+            prefix=[*head, *gap, *state, *newline],
+            options=[[*open_, *option, *close] for option in options],
+            option_read=[len(open_) + len(option) - 1 for option in options],
+            suffix=list(cue),
+            summary=("suffix", -1),
+        )
+
+
+def _delimiters(tokenizer):
+    cached = getattr(tokenizer, "_sysone_qwen_delimiters", None)
+    if cached is None:
+        cached = tuple(tuple(encode_text(tokenizer, text)) for text in ("\n\n", "\n", OPEN, CLOSE, CUE))
+        tokenizer._sysone_qwen_delimiters = cached
+    return cached

sysone/architectures/registry.py

@@ -0,0 +1,53 @@
+"""Look up the architecture for a base model by its `AutoConfig.model_type`."""
+
+from __future__ import annotations
+
+import importlib
+
+from .base import Architecture
+from .modernbert import ModernBertArchitecture
+from .qwen import QwenArchitecture
+
+_REGISTRY: dict[str, type[Architecture]] = {}
+
+
+def register_architecture(model_type: str, cls: type[Architecture]) -> None:
+    """Make `cls` handle base models whose `AutoConfig.model_type` is `model_type`."""
+    if not (isinstance(cls, type) and issubclass(cls, Architecture)):
+        raise TypeError("register_architecture needs an Architecture subclass, got %r" % (cls,))
+    _REGISTRY[model_type] = cls
+
+
+def architecture_for(model_name: str, revision: str | None = None) -> Architecture:
+    from transformers import AutoConfig
+
+    model_type = AutoConfig.from_pretrained(model_name, revision=revision).model_type
+    if model_type not in _REGISTRY:
+        raise ValueError("no architecture is registered for model type %r (from %s); supported: %s. "
+                         "Write an Architecture subclass and call register_architecture(), or pass "
+                         "architecture=" % (model_type, model_name, ", ".join(sorted(_REGISTRY))))
+    return _REGISTRY[model_type]()
+
+
+def architecture_spec(arch: Architecture) -> str:
+    """A string that `architecture_from_spec` turns back into the same architecture."""
+    for model_type, cls in _REGISTRY.items():
+        if type(arch) is cls:
+            return model_type
+    return "%s:%s" % (type(arch).__module__, type(arch).__qualname__)
+
+
+def architecture_from_spec(spec: str) -> Architecture:
+    if spec in _REGISTRY:
+        return _REGISTRY[spec]()
+    module, _, qualname = spec.partition(":")
+    if not qualname:
+        raise ValueError("unknown architecture %r" % spec)
+    obj = importlib.import_module(module)
+    for part in qualname.split("."):
+        obj = getattr(obj, part)
+    return obj()
+
+
+register_architecture("modernbert", ModernBertArchitecture)
+register_architecture("qwen3", QwenArchitecture)