sysone 0.2.0__py3-none-any.whl

sysone/cache.py

@@ -0,0 +1,196 @@
+"""The on-disk feature cache for a frozen backbone (`cache_dir`).
+
+Vectors are grouped into **sets**, one folder per combination of everything that determines them
+(base model and revision, kept layers, layout, precision, prompt format, token limits). Within a
+set, each question is looked up by a hash of its encoded tokens, so growing the data only
+computes the new questions. Each write adds one shard file to the set.
+
+The whole folder is capped at `cache_max_gb`: the least recently used sets are deleted first, and
+a set that would exceed the cap on its own is never written. Nothing is written if it would
+leave less than 10% of the disk free.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import shutil
+import time
+from array import array
+from collections.abc import Callable, Sequence
+from datetime import datetime, timezone
+from pathlib import Path
+
+import torch
+
+from ._util import format_bytes
+from .features import CACHE_DTYPE, Features
+from .layout import Encoded
+
+FORMAT = 1
+SET_PREFIX = "set-"
+META_FILE = "set.json"
+#: Writing is skipped when it would leave less than this fraction of the disk free.
+MIN_FREE_FRACTION = 0.10
+
+
+def question_key(encoded: Encoded) -> bytes:
+    """A 16-byte hash of everything about an encoded question that its vectors depend on."""
+    h = hashlib.blake2b(digest_size=16)
+    for values in (encoded.input_ids, encoded.position_ids, encoded.option_ids, encoded.read_positions,
+                   [encoded.summary_position]):
+        h.update(array("q", values).tobytes())
+        h.update(b"|")
+    return h.digest()
+
+
+def _set_name(parts: dict) -> str:
+    text = json.dumps({"format": FORMAT, **parts}, sort_keys=True, default=str)
+    return SET_PREFIX + hashlib.sha256(text.encode()).hexdigest()[:16]
+
+
+def _folder_bytes(path: Path) -> int:
+    return sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
+
+
+def _read_meta(path: Path) -> dict | None:
+    try:
+        with open(path / META_FILE, encoding="utf-8") as f:
+            return json.load(f)
+    except (OSError, ValueError):
+        return None
+
+
+def _write_json(path: Path, data: dict) -> None:
+    tmp = path.with_suffix(".tmp")
+    with open(tmp, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, default=str)
+    os.replace(tmp, path)
+
+
+def _sets(root: Path) -> list[tuple[Path, dict]]:
+    if not root.is_dir():
+        return []
+    found = []
+    for path in sorted(root.glob(SET_PREFIX + "*")):
+        meta = _read_meta(path) if path.is_dir() else None
+        if meta is not None:
+            found.append((path, meta))
+    return found
+
+
+class FeatureCache:
+    """One set of cached vectors in `root`, identified by `parts`."""
+
+    def __init__(self, root, parts: dict, max_bytes: float, log: Callable[[str], None]):
+        self.root = Path(root).expanduser()
+        self.parts = parts
+        self.path = self.root / _set_name(parts)
+        self.max_bytes = max_bytes
+        self.log = log
+
+    def _touch(self, meta: dict) -> None:
+        meta["last_used"] = time.time()
+        _write_json(self.path / META_FILE, meta)
+
+    def lookup(self, keys: Sequence[bytes]) -> dict[bytes, Features]:
+        """The cached vectors for whichever of `keys` are in this set."""
+        from safetensors.torch import load_file
+
+        meta = _read_meta(self.path)
+        if meta is None:
+            return {}
+        wanted, found = set(keys), {}
+        for shard in sorted(self.path.glob("shard-*.safetensors")):
+            data = load_file(str(shard))
+            offsets = data["offsets"].tolist()
+            for i, key in enumerate(bytes(row) for row in data["keys"].numpy()):
+                if key in wanted and key not in found:
+                    found[key] = Features(data["options"][offsets[i]: offsets[i + 1]].clone(),
+                                          data["summary"][i].clone())
+        self._touch(meta)
+        return found
+
+    def add(self, keys: Sequence[bytes], features: Sequence[Features]) -> int:
+        """Write new vectors as one shard; returns the bytes written (0 when skipped)."""
+        from safetensors.torch import save_file
+
+        if not keys:
+            return 0
+        tensors = {
+            "keys": torch.tensor([list(k) for k in keys], dtype=torch.uint8),
+            "offsets": torch.tensor([0] + [f.options.shape[0] for f in features], dtype=torch.int64).cumsum(0),
+            "options": torch.cat([f.options for f in features]).to(CACHE_DTYPE).contiguous(),
+            "summary": torch.stack([f.summary for f in features]).to(CACHE_DTYPE).contiguous(),
+        }
+        size = sum(t.numel() * t.element_size() for t in tensors.values())
+        current = _folder_bytes(self.path) if self.path.is_dir() else 0
+        if current + size > self.max_bytes:
+            self.log("feature cache: this set would take %s, more than cache_max_gb=%g; not writing it"
+                     % (format_bytes(current + size), self.max_bytes / 1024**3))
+            return 0
+        self.root.mkdir(parents=True, exist_ok=True)
+        disk = shutil.disk_usage(self.root)
+        if disk.free - size < MIN_FREE_FRACTION * disk.total:
+            self.log("feature cache: writing %s would leave less than %d%% of the disk free; not writing it"
+                     % (format_bytes(size), int(100 * MIN_FREE_FRACTION)))
+            return 0
+        self._evict(size)
+        self.path.mkdir(exist_ok=True)
+        meta = _read_meta(self.path) or {"format": FORMAT, "parts": self.parts, "created": time.time(),
+                                         "n_questions": 0}
+        n_shards = len(list(self.path.glob("shard-*.safetensors")))
+        shard = self.path / ("shard-%05d.safetensors" % n_shards)
+        tmp = shard.with_suffix(".tmp")
+        save_file(tensors, str(tmp))
+        os.replace(tmp, shard)
+        meta["n_questions"] += len(keys)
+        self._touch(meta)
+        return shard.stat().st_size
+
+    def _evict(self, incoming: int) -> None:
+        """Delete the least recently used other sets until `incoming` more bytes fit the cap."""
+        others = sorted(((p, m) for p, m in _sets(self.root) if p != self.path),
+                        key=lambda pm: pm[1].get("last_used", 0))
+        total = _folder_bytes(self.root)
+        while others and total + incoming > self.max_bytes:
+            path, _ = others.pop(0)
+            freed = _folder_bytes(path)
+            shutil.rmtree(path, ignore_errors=True)
+            total -= freed
+            self.log("feature cache: deleted the least recently used set %s (%s) to stay under cache_max_gb"
+                     % (path.name, format_bytes(freed)))
+
+    def folder_bytes(self) -> int:
+        return _folder_bytes(self.root) if self.root.is_dir() else 0
+
+
+def cache_info(cache_dir) -> list[dict]:
+    """The cache sets in `cache_dir`, most recently used first: their folder, size, number of
+    questions, last use, and the settings that identify them."""
+    rows = []
+    for path, meta in _sets(Path(cache_dir).expanduser()):
+        size = _folder_bytes(path)
+        last = meta.get("last_used")
+        rows.append({
+            "set": path.name,
+            "path": str(path),
+            "model": meta.get("parts", {}).get("model_name"),
+            "n_questions": meta.get("n_questions", 0),
+            "size_bytes": size,
+            "size": format_bytes(size),
+            "last_used": datetime.fromtimestamp(last, timezone.utc).isoformat(timespec="seconds") if last else None,
+            "parts": meta.get("parts", {}),
+        })
+    return sorted(rows, key=lambda r: r["last_used"] or "", reverse=True)
+
+
+def clear_cache(cache_dir) -> int:
+    """Delete every cache set in `cache_dir` and return the bytes freed. Other files in the folder
+    are left alone."""
+    freed = 0
+    for path, _ in _sets(Path(cache_dir).expanduser()):
+        freed += _folder_bytes(path)
+        shutil.rmtree(path, ignore_errors=True)
+    return freed

sysone/checkpoint.py

@@ -0,0 +1,96 @@
+"""Files of a saved model beyond `config.json` and the head.
+
+- `adapter/`: LoRA adapters in peft's format (`adapter_config.json`, `adapter_model.safetensors`);
+- `backbone.safetensors`: only the fully trained layers and embeddings, in the backbone's precision;
+- `training_state/`: what a later `fit` needs to resume the same run.
+
+Frozen backbone weights are never saved.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from .training import TrainState
+from .tuning import add_adapters, is_adapter
+
+ADAPTER_DIR = "adapter"
+ADAPTER_WEIGHTS = "adapter_model.safetensors"
+BACKBONE_FILE = "backbone.safetensors"
+STATE_DIR = "training_state"
+_STATE_META = "state.json"
+_STATE_TENSORS = "state.pt"
+#: peft stores adapter weights without the adapter's name, under the wrapper's prefix.
+_PEFT_PREFIX = "base_model.model."
+_ADAPTER_KEY = re.compile(r"\.(lora_[AB])\.default\.")
+_PEFT_KEY = re.compile(r"\.(lora_[AB])\.")
+
+
+def save_adapter(backbone: nn.Module, path: Path) -> None:
+    from safetensors.torch import save_file
+
+    path.mkdir(parents=True, exist_ok=True)
+    backbone.sysone_lora_config.save_pretrained(str(path))
+    weights = {_PEFT_PREFIX + _ADAPTER_KEY.sub(r".\1.", name): p.detach().cpu().contiguous()
+               for name, p in backbone.named_parameters() if is_adapter(name)}
+    save_file(weights, str(path / ADAPTER_WEIGHTS))
+
+
+def load_adapter(backbone: nn.Module, path: Path) -> None:
+    from peft import LoraConfig
+    from safetensors.torch import load_file
+
+    add_adapters(backbone, LoraConfig.from_pretrained(str(path)))
+    params = dict(backbone.named_parameters())
+    for key, value in load_file(str(path / ADAPTER_WEIGHTS)).items():
+        name = _PEFT_KEY.sub(r".\1.default.", key.removeprefix(_PEFT_PREFIX))
+        if name not in params:
+            raise ValueError("the saved adapter has a weight %r that this model doesn't" % key)
+        with torch.no_grad():
+            params[name].copy_(value)
+
+
+def save_backbone_weights(named: list[tuple[str, nn.Parameter]], dtype: torch.dtype, path: Path) -> None:
+    from safetensors.torch import save_file
+
+    save_file({name: p.detach().to("cpu", dtype).contiguous() for name, p in named}, str(path))
+
+
+def load_backbone_weights(backbone: nn.Module, path: Path) -> None:
+    from safetensors.torch import load_file
+
+    weights = load_file(str(path))
+    missing = set(weights) - set(dict(backbone.named_parameters()))
+    if missing:
+        raise ValueError("%s has weights this model doesn't: %s" % (path.name, sorted(missing)[:5]))
+    backbone.load_state_dict(weights, strict=False)
+
+
+def save_training_state(state: TrainState, path: Path) -> None:
+    path.mkdir(parents=True, exist_ok=True)
+    meta = {k: getattr(state, k) for k in ("fingerprint", "step", "last_check", "best_step", "running_loss",
+                                           "running_n")}
+    meta["best_loss"] = state.best_loss if math.isfinite(state.best_loss) else None
+    with open(path / _STATE_META, "w", encoding="utf-8") as f:
+        json.dump(meta, f, indent=2)
+    torch.save({"epoch_generator": state.epoch_generator, "rng": state.rng, "optimizer": state.optimizer,
+                "current": state.current}, path / _STATE_TENSORS)
+
+
+def load_training_state(path: Path) -> TrainState | None:
+    if not (path / _STATE_META).is_file():
+        return None
+    with open(path / _STATE_META, encoding="utf-8") as f:
+        meta = json.load(f)
+    tensors = torch.load(path / _STATE_TENSORS, map_location="cpu", weights_only=True)
+    best = meta.pop("best_loss")
+    if "stale" in meta:  # saved before patience was counted in epochs: assume the best check was the last one
+        meta.pop("stale")
+        meta["best_step"] = meta["last_check"]
+    return TrainState(**meta, best_loss=math.inf if best is None else best, **tensors)

sysone/data.py

@@ -0,0 +1,190 @@
+"""Turn the accepted input forms into examples: one per (state, question, label).
+
+Accepted forms:
+
+1. Records: an iterable of dicts with `state`, `questions`, and labels in `expected`, `gold` or a
+   `label` inside each question. Rows may mix question sets freely.
+2. Records plus a question catalog (`questions={id: question_dict}`), so rows carry only `state`
+   and their labels.
+3. A long DataFrame, one row per (state, question): `state | question_id | answer` with a
+   catalog, or self-contained `state | type | instructions | criteria | answer`.
+
+A path to a `.jsonl`, `.json`, `.csv` or `.parquet` file is read first. A `datasets.Dataset` is
+iterated as records. Fields holding JSON text (common in published datasets) are decoded.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections import Counter
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from .questions import Question, parse_question, target_from_gold, target_from_label
+
+_MISSING = object()
+
+
+@dataclass
+class Example:
+    """One question about one state. `target` is None at prediction time."""
+
+    state: str
+    question: Question
+    qid: str
+    target: list[float] | None = None
+    category: str | None = None
+    row: int = 0  # index of the source row; questions from one row share a state
+
+
+def state_text(state: Any) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False, default=str)
+
+
+def _maybe_json(value: Any) -> Any:
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped[:1] in ("{", "["):
+            try:
+                return json.loads(stripped)
+            except json.JSONDecodeError:
+                return value
+    return value
+
+
+def _is_frame(data: Any) -> bool:
+    return hasattr(data, "columns") and hasattr(data, "to_dict") and hasattr(data, "iterrows")
+
+
+def _read_path(path: Path) -> Any:
+    suffix = path.suffix.lower()
+    if suffix == ".jsonl":
+        with open(path, encoding="utf-8") as f:
+            return [json.loads(line) for line in f if line.strip()]
+    if suffix == ".json":
+        with open(path, encoding="utf-8") as f:
+            return json.load(f)
+    if suffix in (".csv", ".parquet"):
+        try:
+            import pandas as pd
+        except ImportError as exc:
+            raise ImportError("reading %s files needs pandas: pip install sysone[pandas]" % suffix) from exc
+        return pd.read_csv(path) if suffix == ".csv" else pd.read_parquet(path)
+    raise ValueError("cannot read %s: use a .jsonl, .json, .csv or .parquet file" % path)
+
+
+def _frame_rows(frame: Any) -> Iterable[dict]:
+    columns = set(frame.columns)
+    if "state" not in columns:
+        raise ValueError("a DataFrame needs a 'state' column")
+    if "question_id" in columns:
+        mode = "catalog"
+    elif {"type", "instructions"} <= columns:
+        mode = "inline"
+    else:
+        raise ValueError("a DataFrame needs either a 'question_id' column (with questions=catalog) or "
+                         "'type' and 'instructions' columns")
+    has_answer = "answer" in columns
+    for record in frame.to_dict(orient="records"):
+        row: dict = {"state": record["state"]}
+        if "category" in columns and record.get("category") is not None:
+            row["category"] = record["category"]
+        if mode == "catalog":
+            qid = str(record["question_id"])
+            row["questions"] = [qid]
+        else:
+            qid = "question"
+            qdef = {"type": record["type"], "instructions": record["instructions"]}
+            criteria = _maybe_json(record.get("criteria"))
+            if criteria is not None and not (isinstance(criteria, float) and criteria != criteria):  # NaN
+                qdef["criteria"] = criteria
+            row["questions"] = {qid: qdef}
+        if has_answer and not _is_null(record["answer"]):
+            row["expected"] = {qid: record["answer"]}
+        yield row
+
+
+def _is_null(value: Any) -> bool:
+    return value is None or (isinstance(value, float) and value != value)
+
+
+def _rows(data: Any) -> Iterable[dict]:
+    if isinstance(data, (str, os.PathLike)):
+        data = _read_path(Path(data))
+    if _is_frame(data):
+        return _frame_rows(data)
+    if isinstance(data, Mapping):
+        raise TypeError("pass a list of rows (dicts), not a single dict")
+    return data
+
+
+def _resolve_questions(row: dict, catalog: Mapping | None, where: str) -> dict:
+    questions = _maybe_json(row.get("questions", _MISSING))
+    if questions is _MISSING or questions is None:
+        if catalog is None:
+            raise ValueError("%s: no 'questions' and no questions= catalog was given" % where)
+        referenced = list(_maybe_json(row.get("expected")) or {}) + list(_maybe_json(row.get("gold")) or {})
+        questions = {qid: catalog[qid] for qid in dict.fromkeys(referenced) if _in_catalog(catalog, qid, where)}
+    elif isinstance(questions, (list, tuple)):
+        # a list of question ids, looked up in the catalog
+        if catalog is None:
+            raise ValueError("%s: 'questions' lists ids but no questions= catalog was given" % where)
+        questions = {qid: catalog[qid] for qid in questions if _in_catalog(catalog, qid, where)}
+    elif not isinstance(questions, dict):
+        raise ValueError("%s: 'questions' must be a dict of question id -> question" % where)
+    return {qid: _maybe_json(qdef) for qid, qdef in questions.items()}
+
+
+def _in_catalog(catalog: Mapping, qid: Any, where: str) -> bool:
+    if qid not in catalog:
+        raise ValueError("%s: question %r is not in the questions= catalog" % (where, qid))
+    return True
+
+
+def examples_from(data: Any, questions: Mapping | None = None, *, labeled: bool = True,
+                  label_smoothing: float = 0.0) -> tuple[list[Example], Counter]:
+    """Examples from any accepted input form.
+
+    With `labeled=True` (training), questions without a label are skipped and counted in the
+    returned `Counter` under "unlabeled". Malformed rows or questions raise `ValueError` naming
+    the row and question.
+    """
+    examples: list[Example] = []
+    skipped: Counter = Counter()
+    for i, row in enumerate(_rows(data)):
+        where = "row %d" % i
+        if not isinstance(row, Mapping):
+            raise ValueError("%s: expected a dict, got %s" % (where, type(row).__name__))
+        if "state" not in row:
+            raise ValueError("%s: missing 'state'" % where)
+        state = state_text(row["state"])
+        category = row.get("category")
+        category = None if _is_null(category) else str(category)
+        qdefs = _resolve_questions(row, questions, where)
+        expected = _maybe_json(row.get("expected")) or {}
+        gold = _maybe_json(row.get("gold")) or {}
+        for qid, qdef in qdefs.items():
+            qwhere = "%s, question %r" % (where, qid)
+            question = parse_question(qid, qdef)
+            target = None
+            if labeled:
+                sources = [name for name, present in (("expected", qid in expected), ("gold", qid in gold),
+                                                      ("label", isinstance(qdef, dict) and "label" in qdef))
+                           if present]
+                if len(sources) > 1:
+                    raise ValueError("%s: has labels in both %s; give exactly one" % (qwhere, " and ".join(sources)))
+                if not sources:
+                    skipped["unlabeled"] += 1
+                    continue
+                if sources[0] == "gold":
+                    target = target_from_gold(question, _maybe_json(gold[qid]), where=qwhere)
+                else:
+                    value = expected[qid] if sources[0] == "expected" else qdef["label"]
+                    target = target_from_label(question, value, label_smoothing, where=qwhere)
+            examples.append(Example(state, question, str(qid), target, category, i))
+    return examples, skipped

sysone/encoding.py

@@ -0,0 +1,80 @@
+"""Turn (state, question) pairs into token layouts, with truncation."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+from .architectures.base import Architecture, encode_text
+from .layout import Encoded, build
+from .prompts import PromptFormat
+from .questions import Question
+
+#: The percentile of training-data lengths that "auto" token limits are set to.
+AUTO_PERCENTILE = 95
+#: Room kept for the question, options and delimiters when capping an "auto" state limit.
+_RESERVED_TOKENS = 256
+_MAX_AUTO_OPTION_TOKENS = 256
+
+
+class Encoder:
+    def __init__(self, arch: Architecture, tokenizer, prompt: PromptFormat, layout: str,
+                 max_state_tokens: int, max_option_tokens: int, max_len: int):
+        self.arch = arch
+        self.tokenizer = tokenizer
+        self.prompt = prompt
+        self.layout = layout
+        self.max_state_tokens = max_state_tokens
+        self.max_option_tokens = max_option_tokens
+        self.max_len = max_len
+        self._cache: dict[str, list[int]] = {}
+
+    def tokens(self, text: str) -> list[int]:
+        ids = self._cache.get(text)
+        if ids is None:
+            ids = encode_text(self.tokenizer, text)
+            self._cache[text] = ids
+        return ids
+
+    def encode(self, state: str, question: Question, order: Sequence[int] | None = None) -> Encoded:
+        """Lay out one question about one state. Slot `s` holds option `order[s]` (default: the
+        question's own order)."""
+        head = self.tokens(self.prompt.head(question))
+        texts = self.prompt.options(question)
+        if order is not None:
+            texts = [texts[i] for i in order]
+        options = [self.tokens(t)[: self.max_option_tokens] for t in texts]
+        state_ids = self.tokens(state)[: self.max_state_tokens]
+        encoded = build(self.arch.pieces(self.tokenizer, head, state_ids, options), self.layout)
+        overflow = len(encoded.input_ids) - self.max_len
+        if overflow > 0:
+            if overflow > len(state_ids):
+                raise ValueError("the question %r with %d options needs more than max_len=%d tokens even "
+                                 "with no state; shorten its options or raise max_len"
+                                 % (question.instructions[:60], question.k, self.max_len))
+            state_ids = state_ids[: len(state_ids) - overflow]
+            encoded = build(self.arch.pieces(self.tokenizer, head, state_ids, options), self.layout)
+        return encoded
+
+
+def _percentile(values: list[int], q: float) -> int:
+    if not values:
+        return 1
+    ordered = sorted(values)
+    index = min(len(ordered) - 1, max(0, math.ceil(q / 100 * len(ordered)) - 1))
+    return max(1, ordered[index])
+
+
+def measure_limits(encoder: Encoder, states: Sequence[str], questions: Sequence[Question],
+                   max_state_tokens, max_option_tokens) -> tuple[int, int]:
+    """Resolve "auto" token limits from the training data: the 95th percentile of token lengths,
+    capped so a sequence still fits the model."""
+    if max_state_tokens == "auto":
+        lengths = [len(encoder.tokens(s)) for s in dict.fromkeys(states)]
+        cap = max(1, encoder.max_len - _RESERVED_TOKENS)
+        max_state_tokens = min(_percentile(lengths, AUTO_PERCENTILE), cap)
+    if max_option_tokens == "auto":
+        texts = {t for q in dict.fromkeys(questions) for t in encoder.prompt.options(q)}
+        lengths = [len(encoder.tokens(t)) for t in texts]
+        max_option_tokens = min(_percentile(lengths, AUTO_PERCENTILE), _MAX_AUTO_OPTION_TOKENS)
+    return int(max_state_tokens), int(max_option_tokens)