sysone 0.2.0__py3-none-any.whl

sysone/tuning.py

@@ -0,0 +1,210 @@
+"""How much of the backbone learns: layer plans, freezing, and LoRA / QLoRA adapters.
+
+The kept layers are split from the top down:
+
+    bottom ── frozen ── adapters (n_lora) ── fully trained (n_full) ── top
+
+The head always trains. With `train_embeddings` the token embeddings train too.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import re
+import sys
+import warnings
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ._util import resolve_count
+from .architectures.base import Architecture
+
+#: (trainable_layers, lora_layers) when the user leaves them at None.
+_DEFAULT_LAYERS = {"head": (0, 0), "qlora": (0, 0.6), "lora": (0, 0.6), "full": ("all", 0)}
+#: Adapter weights carry this in their parameter names (peft's naming).
+_ADAPTER_NAME = re.compile(r"\.lora_[AB]\.")
+
+
+@dataclass(frozen=True)
+class LayerPlan:
+    n_layers: int  # layers kept by num_layers
+    n_full: int  # fully trained top layers
+    n_lora: int  # layers with adapters, directly below the fully trained ones
+    train_embeddings: bool = False
+
+    @property
+    def full_range(self) -> range:
+        return range(self.n_layers - self.n_full, self.n_layers)
+
+    @property
+    def lora_range(self) -> range:
+        top = self.n_layers - self.n_full
+        return range(top - self.n_lora, top)
+
+    @property
+    def learns(self) -> bool:
+        """True when anything in the backbone trains, so its outputs change every step."""
+        return bool(self.n_full or self.n_lora or self.train_embeddings)
+
+    @property
+    def lowest_trained(self) -> int:
+        """The lowest layer gradients flow through (`n_layers` when the backbone is frozen)."""
+        if self.train_embeddings:
+            return 0
+        return self.n_layers - self.n_full - self.n_lora
+
+
+def layer_plan(tuning: str, n_layers: int, trainable_layers, lora_layers, train_embeddings: bool) -> LayerPlan:
+    full_default, lora_default = _DEFAULT_LAYERS[tuning]
+    n_full = resolve_count(full_default if trainable_layers is None else trainable_layers, n_layers,
+                           "trainable_layers")
+    n_lora = resolve_count(lora_default if lora_layers is None else lora_layers, n_layers, "lora_layers")
+    if n_full + n_lora > n_layers:
+        raise ValueError("trainable_layers + lora_layers = %d + %d is more than the %d kept layers"
+                         % (n_full, n_lora, n_layers))
+    return LayerPlan(n_layers, n_full, n_lora, bool(train_embeddings))
+
+
+def quantization_available(device: torch.device) -> bool:
+    """4-bit and 8-bit loading need bitsandbytes and, in practice, an NVIDIA GPU."""
+    if device.type != "cuda":
+        return False
+    try:
+        import bitsandbytes  # noqa: F401
+    except ImportError:
+        return False
+    return True
+
+
+def effective_tuning(tuning: str, device: torch.device, stacklevel: int = 2) -> str:
+    """`tuning`, except that "qlora" trains as "lora" where quantized loading is unavailable."""
+    if tuning == "qlora" and not quantization_available(device):
+        warnings.warn("tuning='qlora' needs bitsandbytes and an NVIDIA GPU, which aren't available here "
+                      "(device=%s); using tuning='lora' instead: the same adapters, with the frozen weights "
+                      "unquantized, so only memory use differs" % device, UserWarning, stacklevel=stacklevel + 1)
+        return "lora"
+    return tuning
+
+
+def quantization_config(bits: int, quant_type: str, compute_dtype: torch.dtype):
+    from transformers import BitsAndBytesConfig
+
+    if bits == 4:
+        return BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type=quant_type,
+                                  bnb_4bit_compute_dtype=compute_dtype)
+    return BitsAndBytesConfig(load_in_8bit=True)
+
+
+def lora_config(arch: Architecture, backbone: nn.Module, plan: LayerPlan, targets, rank: int, alpha, dropout: float):
+    from peft import LoraConfig
+
+    modules = arch.adapter_modules(backbone, targets, plan.lora_range)
+    if not modules:
+        raise ValueError("lora_targets=%r found no linear modules in layers %d-%d"
+                         % (targets, plan.lora_range.start, plan.lora_range.stop - 1))
+    return LoraConfig(r=rank, lora_alpha=2 * rank if alpha is None else alpha, lora_dropout=dropout,
+                      target_modules=modules, bias="none")
+
+
+def add_adapters(backbone: nn.Module, config) -> None:
+    """Put LoRA adapters on `backbone` in place; the model stays a plain Hugging Face model."""
+    from peft import inject_adapter_in_model
+
+    with _ignore_incompatible_torchao():
+        inject_adapter_in_model(config, backbone)
+    backbone.sysone_lora_config = config
+
+
+@contextlib.contextmanager
+def _ignore_incompatible_torchao():
+    """peft refuses to add any adapter when an older torchao is installed (Google Colab ships one),
+    although sysone never uses torchao. In that case, tell peft torchao isn't available while
+    adapters are added, instead of asking the user to uninstall it."""
+    from peft import import_utils
+
+    if not hasattr(import_utils, "is_torchao_available"):  # older peft: no such check
+        yield
+        return
+    try:
+        import_utils.is_torchao_available()
+        problem = None
+    except ImportError as exc:
+        problem = exc
+    if problem is None:
+        yield
+        return
+    warnings.warn("ignoring the installed torchao (%s); sysone doesn't use it" % problem, UserWarning, stacklevel=4)
+    original = import_utils.is_torchao_available
+    holders = [m for name, m in list(sys.modules.items())
+               if name.startswith("peft") and getattr(m, "is_torchao_available", None) is original]
+    for module in holders:
+        module.is_torchao_available = _torchao_unavailable
+    try:
+        yield
+    finally:
+        for module in holders:
+            module.is_torchao_available = original
+
+
+def _torchao_unavailable() -> bool:
+    return False
+
+
+def is_adapter(name: str) -> bool:
+    return bool(_ADAPTER_NAME.search(name))
+
+
+def trained_parameters(arch: Architecture, backbone: nn.Module, plan: LayerPlan) -> dict[str, list]:
+    """`{"adapter": [(name, param)], "full": [(name, param)]}`: the backbone weights that train,
+    in a fixed order. "full" holds fully trained layers and embeddings."""
+    layers = arch.layers(backbone)
+    full = {id(p) for i in plan.full_range for p in layers[i].parameters()}
+    if plan.train_embeddings:
+        full |= {id(p) for p in arch.embeddings(backbone).parameters()}
+    groups: dict[str, list] = {"adapter": [], "full": []}
+    for name, p in backbone.named_parameters():
+        if is_adapter(name):
+            groups["adapter"].append((name, p))
+        elif id(p) in full:
+            groups["full"].append((name, p))
+    return groups
+
+
+def round_trained(groups: dict[str, list], dtype: torch.dtype) -> None:
+    """Store fully trained layers in the backbone's precision once training ends, as they are saved,
+    so a model predicts the same before and after `save` / `load`. Adapters stay in fp32."""
+    for _, p in groups.get("full", []):
+        p.requires_grad_(False)
+        p.data = p.data.to(dtype)
+
+
+def set_trainable(arch: Architecture, backbone: nn.Module, plan: LayerPlan) -> dict[str, list]:
+    """Freeze everything except the planned weights, which are kept in fp32 for the optimizer."""
+    backbone.requires_grad_(False)
+    groups = trained_parameters(arch, backbone, plan)
+    for _, p in groups["adapter"] + groups["full"]:
+        if p.dtype != torch.float32:
+            p.data = p.data.float()
+        p.requires_grad_(True)
+    return groups
+
+
+def wants_checkpointing(setting, plan: LayerPlan, backbone: nn.Module, seq_len: int, micro_batch: int) -> bool:
+    """Whether to turn on gradient checkpointing.
+
+    "auto" turns it on when the activations kept for the backward pass would take more memory
+    than the backbone's weights: about `34·s·h + 5·a·s²` bytes per example and layer in half
+    precision, for the layers gradients flow through.
+    """
+    if setting != "auto":
+        return bool(setting)
+    if not plan.learns:
+        return False
+    config = backbone.config
+    layers = plan.n_layers - plan.lowest_trained
+    per_layer = 34 * seq_len * config.hidden_size + 5 * config.num_attention_heads * seq_len * seq_len
+    activations = micro_batch * layers * per_layer
+    weights = sum(p.numel() * p.element_size() for p in backbone.parameters())
+    return activations > weights

sysone/version.py

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

sysone-0.2.0.dist-info/METADATA

@@ -0,0 +1,144 @@
+Metadata-Version: 2.5
+Name: sysone
+Version: 0.2.0
+Summary: Fine-tune and calibrate system-one decision models with a scikit-learn style fit / predict interface.
+Project-URL: Homepage, https://github.com/ForwardMetric/sysone
+Project-URL: Repository, https://github.com/ForwardMetric/sysone
+Project-URL: Issues, https://github.com/ForwardMetric/sysone/issues
+Project-URL: Design, https://github.com/ForwardMetric/sysone/blob/main/docs/DESIGN.md
+Author: Goutham Balaraman
+License-Expression: MIT
+License-File: LICENSE
+Keywords: calibration,classification,decision models,fine-tuning,lora,modernbert,qlora,qwen,transformers
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: bitsandbytes>=0.45; sys_platform == 'linux' or sys_platform == 'win32'
+Requires-Dist: huggingface-hub>=0.25
+Requires-Dist: numpy>=1.26
+Requires-Dist: peft>=0.13
+Requires-Dist: safetensors>=0.4
+Requires-Dist: torch>=2.4
+Requires-Dist: transformers>=5.0
+Provides-Extra: pandas
+Requires-Dist: pandas>=2.0; extra == 'pandas'
+Description-Content-Type: text/markdown
+
+# sysone
+
+Fine-tune and calibrate "system one" decision models: small transformers that answer typed
+questions (choice, score, yes/no) about a piece of text in a single forward pass. You pick a base
+model, call `fit` on labelled rows and `predict` on new ones. sysone handles tokenization,
+position ids, attention masks, layer freezing, the decision head and the training loop.
+
+Supported base models: **ModernBERT** (`EncoderModel`) and **Qwen3** (`DecoderModel`). Others can
+be added by writing one architecture class.
+
+## Install
+
+```bash
+pip install sysone
+pip install "sysone[pandas]"   # also installs pandas, for the comparison tables of `compare`
+```
+
+On Linux and Windows this also installs bitsandbytes, for 4-bit training (`tuning="qlora"`).
+
+## Quickstart
+
+```python
+from sysone import DecoderModel
+
+questions = {
+    "team":     {"type": "choice", "instructions": "Which team should handle it?",
+                 "criteria": ["sales", "operations", "servicing"]},
+    "severity": {"type": "score", "instructions": "How severe is it?",
+                 "criteria": ["none", "minor", "major", "critical"]},
+    "refund":   {"type": "noul", "instructions": "Is the customer asking for a refund?"},
+}
+
+rows = [
+    {"state": "Hi, my order is stuck in the warehouse and our line is down.",
+     "category": "operations",
+     "questions": questions,
+     "expected": {"team": "operations", "severity": 3, "refund": False}},
+    # ... more rows; rows may mix different question sets
+]
+
+model = DecoderModel(tuning="head").fit(rows)
+
+model.predict("I'd like my money back for the broken unit.", questions)
+# {"answers": {"team": {"type": "choice", "choice": "servicing", "probabilities": {...}, "confidence": ...},
+#              "severity": {"type": "score", "score": 1.2, "legend": {...}, ...},
+#              "refund": {"type": "noul", "noul": 0.97, ...}},
+#  "usage": {"input_tokens": 118, "output_tokens": 0}}
+
+model.save("my-model")
+model = DecoderModel.load("my-model")
+```
+
+To measure the model on held-out labelled rows, and compare it with simple baselines:
+
+```python
+from sysone import baselines, compare
+
+results = baselines(test_rows, train=rows)      # random guess, average answer
+results["trained"] = model.evaluate(test_rows)
+compare(results)                                # accuracy per question type, as a table
+```
+
+Answers follow the Jev `/v1/systemone` response shape. Training data can be records like the
+above, records plus a shared question catalog, a long DataFrame, or a `.jsonl` / `.csv` /
+`.parquet` file. Labels go in `expected`, in a `label` inside each question, or as soft `gold`
+probabilities.
+
+Example notebooks, each comparing accuracy before and after training:
+
+- [quickstart.ipynb](https://github.com/ForwardMetric/sysone/blob/main/examples/quickstart.ipynb): synthetic support tickets, a few minutes;
+- bev-decision-150K, 18 domains, hard labels:
+  [decoder_bev_decision.ipynb](https://github.com/ForwardMetric/sysone/blob/main/examples/decoder_bev_decision.ipynb) (Qwen3) and
+  [encoder_bev_decision.ipynb](https://github.com/ForwardMetric/sysone/blob/main/examples/encoder_bev_decision.ipynb) (ModernBERT);
+- typed-decisions, soft labels from several raters:
+  [decoder_typed_decisions.ipynb](https://github.com/ForwardMetric/sysone/blob/main/examples/decoder_typed_decisions.ipynb) (Qwen3) and
+  [encoder_typed_decisions.ipynb](https://github.com/ForwardMetric/sysone/blob/main/examples/encoder_typed_decisions.ipynb) (ModernBERT).
+
+Notebooks other than the quickstart are named `<model type>_<dataset>.ipynb`.
+
+## Status
+
+This is an early version (0.2). Training covers every `tuning` mode:
+
+- `"head"`: the backbone is frozen, its outputs are computed once and cached (in memory, and on
+  disk with `cache_dir`), and only the decision head trains;
+- `"lora"`, `"qlora"` and `"full"`: adapters or whole layers of the backbone train along with the
+  head. These run on a laptop GPU (Apple MPS) for small datasets; a CUDA GPU is recommended for
+  full-size runs. `"qlora"` needs bitsandbytes and an NVIDIA GPU and falls back to `"lora"`
+  elsewhere, with a warning. It has trained in real 4-bit on a Colab GPU.
+
+`DecoderModel` (Qwen3) defaults to `"qlora"`. `EncoderModel` (ModernBERT) defaults to `"full"`:
+adapters barely moved ModernBERT, while full tuning reached 0.79 test accuracy on typed-decisions.
+
+Training stops at `epochs`, by early stopping, or at `target_loss`. A run that is interrupted
+(Ctrl-C), or saved and loaded, resumes exactly where it stopped. Calibration and `head="full"` are
+designed but not yet implemented.
+
+See [docs/DESIGN.md](https://github.com/ForwardMetric/sysone/blob/main/docs/DESIGN.md) for the full design: every parameter, the layouts,
+architectures, heads, training and checkpoint format.
+
+## Development
+
+```bash
+uv sync                  # or: pip install -e . with the dev dependencies
+uv run pytest            # fast tests on tiny random-weight models, no downloads
+uv run pytest -m slow    # end-to-end tests on the real base models (~1.8 GB download)
+uv run ruff check src tests
+```

sysone-0.2.0.dist-info/RECORD

@@ -0,0 +1,28 @@
+sysone/__init__.py,sha256=dSv6QpRXTyVFbHBtjbO9FBqerurg_df3OA0QkDPNyNA,623
+sysone/_util.py,sha256=7kB_8dq4GKr-3LDMr4vWdGw6oCewf0nfqraJLCXSLTU,3595
+sysone/answers.py,sha256=LuqmHerUsD-7fTiHQd_rZMCU24_O-IoOZUmLcj8w8sk,2023
+sysone/base.py,sha256=RTrLLwHysrCuhJiCzPlp0QQoI4aMc_SBMXKbxN9prNM,44033
+sysone/cache.py,sha256=Pxo74YQ6in55aJg6ejcdFVArLIJ9wGHJTexOo04TrCs,7947
+sysone/checkpoint.py,sha256=_Sv13u4-pWDpdOJu1s6SHDKprQBqHXZNQWzL5BCjEF4,3897
+sysone/data.py,sha256=mGuUfOkCdbmJqalunaiR3xM1CkcNOMO_1kH-zOPIB5c,8002
+sysone/encoding.py,sha256=OALfxVoDqiwnXfg0Sjgaf6dqRwAsYDn3ZoMOL-Ti7fc,3628
+sysone/evaluation.py,sha256=hICp7VzkRmNT_FW04nnV-fy7Or-rhz_nsMY9FWGWR60,7718
+sysone/features.py,sha256=kgod_IvBxRCkwkY0RIbk8BYMH56qEsuO05OfTcO44Fs,4024
+sysone/heads.py,sha256=FeTWnYvKiQO1bdm3ZAyCxJAZHiWl0bgX2LgVcXC9Tcg,3070
+sysone/layout.py,sha256=tDW40SYMSGpT4JbsXqrGEV7xwxsppQBlGnNd_QsQKyo,5473
+sysone/losses.py,sha256=58gubFLKxZN8S0hM5xyFu3DY50Jup7dKFvlvufVAmQQ,626
+sysone/models.py,sha256=OmA25u1iLNPxKx4EOHUX6JsMeH4weXopOIIsXWat62A,961
+sysone/prompts.py,sha256=0UjA-QpVqePrh1zuvbqlvTkW1P9ZuzA7cFo5NuyZG5I,2190
+sysone/questions.py,sha256=HfgRQCiALD_hG0YMjb39Ic0Cbc5BDeLROLejusK-iQw,8252
+sysone/training.py,sha256=w59lLbjGzKlQ38EyhLYgdv8eNEnbetgSdVdZdBYB144,18043
+sysone/tuning.py,sha256=dZApkJG1SVazXNCDkl0heo_73qDgEaFG0aIS8jq1ezc,8457
+sysone/version.py,sha256=Zn1KFblwuFHiDRdRAiRnDBRkbPttWh44jKa5zG2ov0E,22
+sysone/architectures/__init__.py,sha256=WwbBxetVcw9lsFsvGPHkomTOaS30e2JGp2HUE11mIKQ,317
+sysone/architectures/base.py,sha256=dzMKKb8ALf5ghp8dYCGdGNXkN5Jk6V556tknwV_yEes,7359
+sysone/architectures/modernbert.py,sha256=AoQbZxRjQ2eChR8rlZ6tjVa8_MYApwJhPi6BM-6bSJk,1708
+sysone/architectures/qwen.py,sha256=--tfHT140dQXEmvKBm6ek9gAaoMM15dK0_kjX6zUHMM,1924
+sysone/architectures/registry.py,sha256=7W_hr-hHX3JRKhEBsZLdVh836_C7FqHCm_CXoXre6HI,2044
+sysone-0.2.0.dist-info/METADATA,sha256=P6urB5koPoo7uEGvxVcPnekBjXGOKOzKlgTFzAAdgGg,6714
+sysone-0.2.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+sysone-0.2.0.dist-info/licenses/LICENSE,sha256=vc1flUELWk-ba73Rv_DNFquov4HjzJGku7jzHqSBRiI,1074
+sysone-0.2.0.dist-info/RECORD,,

sysone-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

sysone-0.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Goutham Balaraman
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.