sysone 0.2.0__py3-none-any.whl

sysone/evaluation.py

@@ -0,0 +1,167 @@
+"""Evaluation: score predictions against labels, baselines, and side-by-side comparison.
+
+Labels are read exactly as `fit` reads them (`expected`, a `label` inside a question, or soft
+`gold`), and each becomes a target distribution over the question's options. So one set of
+metrics covers hard and soft labels alike:
+
+- **accuracy**: the predicted top option is a top option of the target. For a hard label this is
+  ordinary accuracy; for soft labels it is agreement with the raters' majority. When several
+  predicted options tie for the top, the share of them that are right counts, so a uniform guess
+  scores exactly 1/k.
+- **distance**: total variation distance between predicted and target distributions
+  (0 = identical, 1 = completely different).
+- **log_loss**: cross-entropy of the predictions against the target.
+- **score_error**: for score questions, the absolute gap between predicted and target expected level.
+- **auc**: for noul questions, how well P(true) ranks true above false targets.
+"""
+
+from __future__ import annotations
+
+import math
+from collections import Counter, defaultdict
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from .data import Example, examples_from
+from .questions import QUESTION_TYPES
+
+METRICS = ("accuracy", "distance", "log_loss", "score_error", "auc")
+_TIE = 1e-9
+_LOG_FLOOR = 1e-12
+
+
+def question_key(example: Example) -> str:
+    """How questions are grouped in `by_question`: `category/qid`, or `qid` without a category."""
+    return "%s/%s" % (example.category, example.qid) if example.category is not None else example.qid
+
+
+def _auc(scores: Sequence[float], positives: Sequence[bool]) -> float | None:
+    pos = [s for s, y in zip(scores, positives) if y]
+    neg = [s for s, y in zip(scores, positives) if not y]
+    if not pos or not neg:
+        return None
+    wins = sum((p > n) + 0.5 * (p == n) for p in pos for n in neg)
+    return wins / (len(pos) * len(neg))
+
+
+def _summarize(items: list[dict]) -> dict:
+    out: dict[str, Any] = {"n": len(items)}
+    out["accuracy"] = sum(it["hit"] for it in items) / len(items)
+    out["distance"] = sum(it["distance"] for it in items) / len(items)
+    out["log_loss"] = sum(it["log_loss"] for it in items) / len(items)
+    gaps = [it["level_gap"] for it in items if it["level_gap"] is not None]
+    if gaps:
+        out["score_error"] = sum(gaps) / len(gaps)
+    noul = [it for it in items if it["type"] == "noul" and it["noul_label"] is not None]
+    auc = _auc([it["p_true"] for it in noul], [it["noul_label"] for it in noul])
+    if auc is not None:
+        out["auc"] = auc
+    return out
+
+
+def score(examples: Sequence[Example], probabilities: Sequence[Sequence[float]],
+          skipped: Mapping | None = None) -> dict:
+    """Score each example's predicted option probabilities against its target.
+
+    Returns a plain dict: `n_questions`, `overall`, `by_type`, `by_category`, `by_question`, and
+    `skipped` (questions without labels). Each group holds `n` and the metrics that apply to it.
+    """
+    items = []
+    for ex, p in zip(examples, probabilities):
+        t = ex.target
+        p = [float(v) for v in p[: len(t)]]
+        best = max(t)
+        # The prediction's top options; if several tie, credit the share of them that are right,
+        # which is the expected accuracy of breaking the tie at random (so a uniform guess scores 1/k).
+        top = max(p)
+        picked = [i for i, v in enumerate(p) if v >= top - _TIE]
+        item = {
+            "type": ex.question.type,
+            "category": ex.category,
+            "question": question_key(ex),
+            "hit": sum(t[i] >= best - _TIE for i in picked) / len(picked),
+            "distance": 0.5 * sum(abs(a - b) for a, b in zip(p, t)),
+            "log_loss": -sum(b * math.log(max(a, _LOG_FLOOR)) for a, b in zip(p, t) if b > 0),
+            "level_gap": None,
+            "p_true": None,
+            "noul_label": None,
+        }
+        if ex.question.type == "score":
+            item["level_gap"] = abs(sum(i * a for i, a in enumerate(p)) - sum(i * b for i, b in enumerate(t)))
+        if ex.question.type == "noul":
+            item["p_true"] = p[1]
+            if abs(t[1] - t[0]) > _TIE:  # a tied soft label has no side
+                item["noul_label"] = t[1] > t[0]
+        items.append(item)
+
+    report: dict[str, Any] = {"n_questions": len(items), "skipped": dict(skipped or {})}
+    if not items:
+        report.update(overall=None, by_type={}, by_category={}, by_question={})
+        return report
+    report["overall"] = _summarize(items)
+    for name, field in (("by_type", "type"), ("by_category", "category"), ("by_question", "question")):
+        groups: dict = defaultdict(list)
+        for it in items:
+            if it[field] is not None:
+                groups[it[field]].append(it)
+        order = QUESTION_TYPES if field == "type" else sorted(groups, key=str)
+        report[name] = {key: _summarize(groups[key]) for key in order if key in groups}
+    return report
+
+
+def baselines(data, train=None, questions: Mapping | None = None) -> dict[str, dict]:
+    """Reports for simple guessers on the labelled `data`, to compare a model against.
+
+    - **random guess**: every option equally likely.
+    - **average answer** (when `train` is given): for each question (`category/qid`, same number
+      of options), the average target over the training data. Where a question never appears in
+      `train`, it falls back to the random guess.
+    """
+    examples, skipped = examples_from(data, questions, labeled=True)
+    uniform = [[1.0 / ex.question.k] * ex.question.k for ex in examples]
+    out = {"random guess": score(examples, uniform, skipped)}
+    if train is not None:
+        train_examples, _ = examples_from(train, questions, labeled=True)
+        sums: dict = {}
+        counts: Counter = Counter()
+        for ex in train_examples:
+            key = (question_key(ex), ex.question.k)
+            acc = sums.setdefault(key, [0.0] * ex.question.k)
+            for i, v in enumerate(ex.target):
+                acc[i] += v
+            counts[key] += 1
+        prior = []
+        for ex in examples:
+            key = (question_key(ex), ex.question.k)
+            prior.append([v / counts[key] for v in sums[key]] if key in sums else [1.0 / ex.question.k] * ex.question.k)
+        out["average answer"] = score(examples, prior, skipped)
+    return out
+
+
+def compare(reports: Mapping[str, dict], by: str = "type", metric: str = "accuracy"):
+    """A pandas table comparing reports side by side, one row per report.
+
+    - `by="overall"`: every metric as a column.
+    - `by="type"`, `"category"` or `"question"`: one column per group, plus `overall`, showing
+      `metric` (default `"accuracy"`).
+    """
+    try:
+        import pandas as pd
+    except ImportError as exc:
+        raise ImportError("compare() needs pandas: pip install sysone[pandas]") from exc
+    if by == "overall":
+        rows = {name: {m: r["overall"].get(m) for m in METRICS if r["overall"] and m in r["overall"]}
+                for name, r in reports.items()}
+        return pd.DataFrame(rows).T
+    key = {"type": "by_type", "category": "by_category", "question": "by_question"}.get(by)
+    if key is None:
+        raise ValueError("by must be 'overall', 'type', 'category' or 'question', got %r" % (by,))
+    if metric not in METRICS:
+        raise ValueError("metric must be one of %s, got %r" % (", ".join(METRICS), metric))
+    rows = {}
+    for name, r in reports.items():
+        row = {group: stats.get(metric) for group, stats in r[key].items()}
+        if r["overall"]:
+            row["overall"] = r["overall"].get(metric)
+        rows[name] = row
+    return pd.DataFrame(rows).T

sysone/features.py

@@ -0,0 +1,98 @@
+"""Option and summary vectors from a frozen backbone.
+
+With `tuning="head"` nothing in the backbone learns, so each example's vectors never change: they
+are computed once and cached on the CPU, and every epoch trains the head on the cache.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+import torch
+
+from ._util import autocast, empty_device_cache, is_out_of_memory
+from .architectures.base import Architecture
+from .layout import Encoded, collate
+
+#: Vectors are cached in bf16: the same exponent range as fp32 at half the memory.
+CACHE_DTYPE = torch.bfloat16
+_START_BATCH = 32
+
+
+@dataclass
+class Features:
+    options: torch.Tensor  # [k, H]
+    summary: torch.Tensor  # [H]
+
+
+def read_vectors(hidden: torch.Tensor, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather option vectors `[B, K, H]` and summary vectors `[B, H]` from hidden states."""
+    rows = torch.arange(hidden.shape[0], device=hidden.device)
+    options = hidden[rows[:, None], batch["read_positions"]]
+    summary = hidden[rows, batch["summary_positions"]]
+    return options, summary
+
+
+def run_in_batches(items: Sequence, fn: Callable[[list], list], start: int = _START_BATCH,
+                   length: Callable | None = None) -> list:
+    """Apply `fn` to batches of `items` and return results in input order.
+
+    Items are batched by similar length, so little compute goes to padding. If a batch runs out of
+    memory, the batch size is halved and the batch retried.
+    """
+    order = sorted(range(len(items)), key=(lambda i: -length(items[i])) if length else None)
+    results: list = [None] * len(items)
+    size, i = start, 0
+    while i < len(order):
+        chunk = order[i: i + size]
+        try:
+            outputs = fn([items[j] for j in chunk])
+        except (RuntimeError, torch.OutOfMemoryError) as exc:
+            if not is_out_of_memory(exc) or size == 1:
+                raise
+            size = max(1, size // 2)
+            empty_device_cache()
+            continue
+        for j, out in zip(chunk, outputs):
+            results[j] = out
+        i += len(chunk)
+    return results
+
+
+@torch.no_grad()
+def extract(arch: Architecture, backbone, encoded: Sequence[Encoded], pad_id: int,
+            device: torch.device, autocast_dtype: torch.dtype | None = None) -> list[Features]:
+    """Run the backbone, without gradients, over `encoded` and return each example's vectors on
+    the CPU. `autocast_dtype` is needed when trained weights are kept in fp32 next to frozen
+    half-precision ones."""
+    backbone.eval()
+
+    def batch_fn(chunk: list[Encoded]) -> list[Features]:
+        batch = {k: v.to(device) for k, v in collate(chunk, pad_id).items()}
+        with autocast(device, autocast_dtype):
+            hidden = arch.hidden_states(backbone, batch)
+        options, summary = read_vectors(hidden, batch)
+        options, summary = options.to("cpu", CACHE_DTYPE), summary.to("cpu", CACHE_DTYPE)
+        return [Features(options[b, : len(e.read_positions)].clone(), summary[b].clone())
+                for b, e in enumerate(chunk)]
+
+    return run_in_batches(encoded, batch_fn, length=lambda e: len(e.input_ids))
+
+
+def collate_features(items: Sequence[Features], device: torch.device) -> tuple[torch.Tensor, ...]:
+    """Pad cached vectors into `(options [B, K, H], option_mask [B, K], summary [B, H])`."""
+    n = len(items)
+    k = max(f.options.shape[0] for f in items)
+    hidden = items[0].summary.shape[-1]
+    options = torch.zeros(n, k, hidden, dtype=torch.float32)
+    mask = torch.zeros(n, k, dtype=torch.bool)
+    for i, f in enumerate(items):
+        options[i, : f.options.shape[0]] = f.options.float()
+        mask[i, : f.options.shape[0]] = True
+    summary = torch.stack([f.summary.float() for f in items])
+    return options.to(device), mask.to(device), summary.to(device)
+
+
+def cache_bytes(n_vectors: int, hidden: int) -> int:
+    return n_vectors * hidden * torch.finfo(CACHE_DTYPE).bits // 8

sysone/heads.py

@@ -0,0 +1,66 @@
+"""Decision heads: turn option vectors and a summary vector into one logit per option."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from .questions import QUESTION_TYPES
+
+#: Logit for padded option slots. Finite, so a zero target times this stays zero in the loss.
+MASKED_LOGIT = -1e4
+
+
+def _num_heads(width: int) -> int:
+    heads = max(1, width // 64)
+    while width % heads:
+        heads -= 1
+    return heads
+
+
+class CompactHead(nn.Module):
+    """Works on `[type embedding, option vectors…, summary vector]`, not the whole sequence.
+
+    The vectors are projected to `head_dim`, passed through `head_layers` pre-norm self-attention
+    blocks with no positional encoding (so the options are treated as a set), and each option is
+    scored as the scaled dot product of a projection of the summary and a projection of the
+    option. This works for any number of options. The head always runs in fp32.
+    """
+
+    def __init__(self, hidden_size: int, head_dim: int, head_layers: int = 2, dropout: float = 0.0):
+        super().__init__()
+        self.head_dim = head_dim
+        self.input_proj = nn.Linear(hidden_size, head_dim)
+        self.type_embedding = nn.Embedding(len(QUESTION_TYPES), head_dim)
+        self.input_norm = nn.LayerNorm(head_dim)
+        if head_layers > 0:
+            layer = nn.TransformerEncoderLayer(head_dim, _num_heads(head_dim), 4 * head_dim, dropout,
+                                               activation="gelu", batch_first=True, norm_first=True)
+            self.blocks = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False)
+        else:
+            self.blocks = None
+        self.final_norm = nn.LayerNorm(head_dim)
+        self.summary_proj = nn.Linear(head_dim, head_dim)
+        self.option_proj = nn.Linear(head_dim, head_dim)
+
+    def forward(self, options: torch.Tensor, option_mask: torch.Tensor, summary: torch.Tensor,
+                qtype: torch.Tensor) -> torch.Tensor:
+        """options `[B, K, H]`, option_mask `[B, K]` (True = real), summary `[B, H]`, qtype `[B]`
+        -> logits `[B, K]`, with padded slots at `MASKED_LOGIT`."""
+        with torch.autocast(device_type=options.device.type, enabled=False):
+            options, summary = options.float(), summary.float()
+            n = options.shape[0]
+            x = torch.cat([self.type_embedding(qtype)[:, None], self.input_proj(options),
+                           self.input_proj(summary)[:, None]], dim=1)
+            x = self.input_norm(x)
+            always = torch.zeros(n, 1, dtype=torch.bool, device=options.device)
+            padding = torch.cat([always, ~option_mask, always], dim=1)  # True = ignore
+            if self.blocks is not None:
+                x = self.blocks(x, src_key_padding_mask=padding)
+            x = self.final_norm(x)
+            s = self.summary_proj(x[:, -1])
+            o = self.option_proj(x[:, 1:-1])
+            logits = torch.einsum("bd,bkd->bk", s, o) / math.sqrt(self.head_dim)
+            return logits.masked_fill(~option_mask, MASKED_LOGIT)

sysone/layout.py

@@ -0,0 +1,129 @@
+"""Token layouts and attention masks shared by every architecture.
+
+An architecture describes a sequence as `Pieces`: shared tokens before the options (`prefix`),
+one token list per option, and shared tokens after them (`suffix`). `build` lays the pieces out:
+
+- **parallel**: every option starts at the same position id, the first after the prefix. Tokens
+  after the options continue from the position after the longest option. With `attention_masks`
+  blocking option-to-option attention, reordering the options only reorders their vectors.
+- **sequential**: ordinary positions and full attention between options.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+LAYOUTS = ("parallel", "sequential")
+
+
+@dataclass
+class Pieces:
+    prefix: list[int]
+    options: list[list[int]]
+    option_read: list[int]  # index, within each option's tokens, of the token its vector is read from
+    suffix: list[int]
+    summary: tuple[str, int]  # ("prefix" | "suffix", index into that piece; negative counts from the end)
+
+
+@dataclass
+class Encoded:
+    input_ids: list[int]
+    position_ids: list[int]
+    option_ids: list[int]  # 0 for shared tokens, s + 1 for tokens of the option in slot s
+    read_positions: list[int]
+    summary_position: int
+
+
+def build(pieces: Pieces, layout: str) -> Encoded:
+    if layout not in LAYOUTS:
+        raise ValueError("layout must be one of %s, got %r" % (", ".join(LAYOUTS), layout))
+    if not pieces.options:
+        raise ValueError("a question needs at least one option")
+    ids = list(pieces.prefix)
+    start = len(ids)
+    positions = list(range(start))
+    option_ids = [0] * start
+    read = []
+    for slot, (tokens, offset) in enumerate(zip(pieces.options, pieces.option_read)):
+        if not tokens:
+            raise ValueError("option %d has no tokens" % slot)
+        read.append(len(ids) + (offset % len(tokens)))
+        if layout == "parallel":
+            positions.extend(range(start, start + len(tokens)))
+        else:
+            positions.extend(range(len(ids), len(ids) + len(tokens)))
+        option_ids.extend([slot + 1 if layout == "parallel" else 0] * len(tokens))
+        ids.extend(tokens)
+    after = positions[-1] + 1 if layout == "sequential" else start + max(len(t) for t in pieces.options)
+    suffix_start = len(ids)
+    ids.extend(pieces.suffix)
+    positions.extend(range(after, after + len(pieces.suffix)))
+    option_ids.extend([0] * len(pieces.suffix))
+
+    where, index = pieces.summary
+    if where == "prefix":
+        summary = index % len(pieces.prefix)
+    elif where == "suffix":
+        summary = suffix_start + index % len(pieces.suffix)
+    else:
+        raise ValueError("summary must be in the 'prefix' or the 'suffix', got %r" % (where,))
+    return Encoded(ids, positions, option_ids, read, summary)
+
+
+def collate(encoded: list[Encoded], pad_id: int) -> dict[str, torch.Tensor]:
+    """Pad a list of encodings into batch tensors."""
+    n = len(encoded)
+    length = max(len(e.input_ids) for e in encoded)
+    k = max(len(e.read_positions) for e in encoded)
+    input_ids = torch.full((n, length), pad_id, dtype=torch.long)
+    position_ids = torch.zeros((n, length), dtype=torch.long)
+    option_ids = torch.zeros((n, length), dtype=torch.long)
+    real = torch.zeros((n, length), dtype=torch.bool)
+    read = torch.zeros((n, k), dtype=torch.long)
+    option_mask = torch.zeros((n, k), dtype=torch.bool)
+    summary = torch.zeros(n, dtype=torch.long)
+    for i, e in enumerate(encoded):
+        m = len(e.input_ids)
+        input_ids[i, :m] = torch.tensor(e.input_ids)
+        position_ids[i, :m] = torch.tensor(e.position_ids)
+        option_ids[i, :m] = torch.tensor(e.option_ids)
+        real[i, :m] = True
+        read[i, : len(e.read_positions)] = torch.tensor(e.read_positions)
+        option_mask[i, : len(e.read_positions)] = True
+        summary[i] = e.summary_position
+    return {
+        "input_ids": input_ids,
+        "position_ids": position_ids,
+        "option_ids": option_ids,
+        "real": real,
+        "read_positions": read,
+        "option_mask": option_mask,
+        "summary_positions": summary,
+    }
+
+
+def attention_mask(real: torch.Tensor, position_ids: torch.Tensor, option_ids: torch.Tensor, *,
+                   causal: bool, sliding_window: int | None = None) -> torch.Tensor:
+    """A boolean `[batch, 1, query, key]` mask, True where attention is allowed.
+
+    A token may attend to every real token, except:
+    - tokens of a different option;
+    - later tokens, when `causal`;
+    - tokens more than `sliding_window` positions away, measured in position ids (options share
+      positions, so a sequence-index window would let one option see more context than another).
+
+    Padded query rows attend only to themselves, so no row is fully masked.
+    """
+    n, length = real.shape
+    allowed = real[:, None, :].expand(n, length, length)
+    if causal:
+        allowed = allowed & torch.ones(length, length, dtype=torch.bool, device=real.device).tril()
+    q_opt, k_opt = option_ids[:, :, None], option_ids[:, None, :]
+    allowed = allowed & ~((q_opt > 0) & (k_opt > 0) & (q_opt != k_opt))
+    if sliding_window is not None:
+        allowed = allowed & ((position_ids[:, :, None] - position_ids[:, None, :]).abs() <= sliding_window)
+    eye = torch.eye(length, dtype=torch.bool, device=real.device)
+    allowed = (allowed & real[:, :, None]) | (eye & ~real[:, :, None])
+    return allowed[:, None]

sysone/losses.py

@@ -0,0 +1,17 @@
+"""Training objectives over each row's own options."""
+
+from __future__ import annotations
+
+import torch
+
+LOSSES = ("soft-ce",)
+
+
+def soft_ce(logits: torch.Tensor, target: torch.Tensor, option_mask: torch.Tensor,
+            reduction: str = "mean") -> torch.Tensor:
+    """Cross-entropy against a target distribution. For one-hot targets this is ordinary
+    cross-entropy. Padded option slots carry no target mass. `reduction="none"` gives one loss
+    per row."""
+    logp = torch.log_softmax(logits.float(), dim=-1)
+    rows = -(target * logp * option_mask).sum(-1)
+    return rows if reduction == "none" else rows.mean()

sysone/models.py

@@ -0,0 +1,24 @@
+"""The public model classes. They share every parameter; only a few defaults differ."""
+
+from __future__ import annotations
+
+from .base import SystemOneModel
+
+
+class EncoderModel(SystemOneModel):
+    """A decision model on a bidirectional encoder, ModernBERT by default.
+
+    Trains every layer by default (`tuning="full"`): adapters barely moved ModernBERT, which was
+    pretrained only to fill in masked words, and a model this size is affordable to train fully.
+    """
+
+    def __init__(self, model_name: str = "answerdotai/ModernBERT-base", *, num_layers="all", tuning="full",
+                 **kwargs):
+        super().__init__(model_name, num_layers=num_layers, tuning=tuning, **kwargs)
+
+
+class DecoderModel(SystemOneModel):
+    """A decision model on a causal decoder, Qwen3 by default."""
+
+    def __init__(self, model_name: str = "Qwen/Qwen3-0.6B", *, num_layers=0.7, **kwargs):
+        super().__init__(model_name, num_layers=num_layers, **kwargs)

sysone/prompts.py

@@ -0,0 +1,60 @@
+"""Prompt formats: the text of each piece of a question.
+
+A model trained on one wording reads another as a different input, so the format is saved with
+the model and reused at prediction time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+
+from .questions import Question
+
+
+def _default_task_lines() -> dict:
+    return {
+        "choice": "Choose the best option.",
+        "score": "Rate it on the given scale.",
+        "noul": "Answer yes or no.",
+    }
+
+
+@dataclass
+class PromptFormat:
+    task_lines: dict = field(default_factory=_default_task_lines)
+    question: str = "Question: {instructions}"
+    choice_option: str = "{label}: {description}"
+    score_option: str = "{index}: {level}"
+    noul_labels: tuple = ("No", "Yes")
+
+    def head(self, question: Question) -> str:
+        """The task line and the question, as one piece of text."""
+        return "%s\n%s" % (self.task_lines[question.type], self.question.format(instructions=question.instructions))
+
+    def options(self, question: Question) -> list[str]:
+        """The text of each option, in the question's option order."""
+        if question.type == "choice":
+            return [
+                str(key) if text is None else self.choice_option.format(label=key, description=text)
+                for key, text in zip(question.keys, question.texts)
+            ]
+        if question.type == "score":
+            # The index must be in the text: in the parallel layout every option shares position
+            # ids, so the text is the only place the order of the levels survives.
+            return [self.score_option.format(index=i, level=text) for i, text in enumerate(question.texts)]
+        return [
+            label if text is None else "%s: %s" % (label, text)
+            for label, text in zip(self.noul_labels, question.texts)
+        ]
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["noul_labels"] = list(self.noul_labels)
+        return d
+
+    @classmethod
+    def from_dict(cls, d: dict) -> PromptFormat:
+        d = dict(d)
+        if "noul_labels" in d:
+            d["noul_labels"] = tuple(d["noul_labels"])
+        return cls(**d)