From e45a5c36f9d0bca3fe9e277a095e17730a806019 Mon Sep 17 00:00:00 2001 From: rschumann Date: Tue, 21 Jul 2026 15:44:55 +0200 Subject: [PATCH] feat: add safetensors checkpoints --- AGENTS.md | 2 +- CHANGELOG.md | 7 + DOCUMENTATION.md | 16 +- README.md | 18 +- examples/config.toml | 2 +- examples/santander_retrieval_demo.ipynb | 26 +- linear_adapter_trainer/adapter/model.py | 218 +++++++++++++-- linear_adapter_trainer/cli.py | 13 +- pyproject.toml | 1 + tests/test_adapter_model.py | 354 +++++++++++++++++++++++- tests/test_cli.py | 32 +++ 11 files changed, 641 insertions(+), 48 deletions(-) create mode 100644 tests/test_cli.py diff --git a/AGENTS.md b/AGENTS.md index dde5439..4b31013 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,7 +15,7 @@ code is public; keep it explicit, typed, and easy to read. - Keep defaults deterministic and seeded. - Avoid hidden behavior and side effects. - Add a dependency only when it removes real complexity. Core stays light - (`numpy`, `torch`, `tqdm`); heavy backends are optional extras. + (`numpy`, `torch`, `safetensors`, `tqdm`); heavy backends are optional extras. - The pipeline must run fully offline via `HashingEmbedder` + `TemplateQueryGenerator` so CI never needs network or API keys. diff --git a/CHANGELOG.md b/CHANGELOG.md index 22c3014..e13a40d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] ### Added +- `LinearAdapter.migrate_checkpoint()` for converting legacy `.pt` and `.pth` + checkpoints to safetensors. - Santander demo notebook (`examples/santander_retrieval_demo.ipynb`): scrape a real corpus, generate a QA dataset, train the adapter, and measure the retrieval improvement end-to-end. @@ -17,6 +19,11 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). - Offline unit tests for the scraper's HTML cleaning. ### Changed +- **Breaking:** `LinearAdapter.save()` now writes only pickle-free + `.safetensors` checkpoints with versioned, strictly validated JSON metadata. + Legacy `.pt` and `.pth` checkpoints remain read-only through + `torch.load(..., weights_only=True)`. The CLI default output path now ends in + `.safetensors`, and configured legacy output extensions fail before training. - Export `LLMQueryGenerator` from the top-level package and de-duplicate `__all__`. diff --git a/DOCUMENTATION.md b/DOCUMENTATION.md index e12b5dc..b4c3161 100644 --- a/DOCUMENTATION.md +++ b/DOCUMENTATION.md @@ -149,9 +149,23 @@ adapter = LinearAdapter(AdapterConfig( normalize_output=True, # L2-normalize outputs )) adapted = adapter.transform(query_matrix) # numpy in/out for inference -adapter.save("adapter.pt"); LinearAdapter.load("adapter.pt") +adapter.save("adapter.safetensors"); LinearAdapter.load("adapter.safetensors") ``` +Safetensors checkpoints contain the adapter tensors plus a format version and +strictly validated JSON configuration in the file metadata. New writes must +use the `.safetensors` extension. Legacy `.pt` and `.pth` checkpoints are read +only through `torch.load(..., weights_only=True)` and emit a deprecation +warning. Convert them with: + +```python +LinearAdapter.migrate_checkpoint("adapter.pt", "adapter.safetensors") +``` + +`load(..., map_location=...)` accepts a device string or `torch.device` for +safetensors checkpoints. Parser selection is extension-based; corrupt +safetensors files never fall back to the legacy loader. + A residual adapter starts as a no-op (`W = 0`), so early training never hurts the base embeddings. A non-residual square adapter is initialized at identity. diff --git a/README.md b/README.md index defe622..9366f99 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ pip install "linear-adapter-trainer[openai]" # OpenAI API pip install "linear-adapter-trainer[all]" ``` -The core install is dependency-light (`numpy`, `torch`, `tqdm`). A +The core install is dependency-light (`numpy`, `torch`, `safetensors`, `tqdm`). A dependency-free `HashingEmbedder` and `TemplateQueryGenerator` let you run the whole pipeline offline (great for CI and demos). @@ -116,7 +116,7 @@ dataset = DatasetGenerator( result = AdapterTrainer(kb, embedder, TrainingConfig(epochs=30)).fit(dataset) print(result.improvement) # delta per metric vs the base embeddings -result.adapter.save("adapter.pt") +result.adapter.save("adapter.safetensors") ``` At query time: @@ -125,11 +125,23 @@ At query time: import numpy as np from linear_adapter_trainer import LinearAdapter -adapter = LinearAdapter.load("adapter.pt") +adapter = LinearAdapter.load("adapter.safetensors") query_vec = embedder.embed(["how do plants make energy?"]) adapted = adapter.transform(query_vec) # use this for nearest-neighbor search ``` +New checkpoints use safetensors with versioned JSON metadata, so saving and +loading them does not use pickle. Existing `.pt` and `.pth` checkpoints remain +readable through PyTorch's restricted `weights_only=True` loader and emit a +deprecation warning. Migrate one explicitly: + +```python +LinearAdapter.migrate_checkpoint("adapter.pt", "adapter.safetensors") +``` + +Checkpoint parsers are selected by extension. A malformed `.safetensors` file +is rejected and is never retried as a legacy PyTorch checkpoint. + ## Quickstart (CLI) Everything is driven by one TOML file (see `examples/config.toml`): diff --git a/examples/config.toml b/examples/config.toml index 21c6517..fef26f0 100644 --- a/examples/config.toml +++ b/examples/config.toml @@ -59,5 +59,5 @@ seed = 0 [output] dataset_dir = "examples/artifacts/dataset" -adapter_path = "examples/artifacts/adapter.pt" +adapter_path = "examples/artifacts/adapter.safetensors" metrics_path = "examples/artifacts/metrics.json" diff --git a/examples/santander_retrieval_demo.ipynb b/examples/santander_retrieval_demo.ipynb index 0d29d58..76c765e 100644 --- a/examples/santander_retrieval_demo.ipynb +++ b/examples/santander_retrieval_demo.ipynb @@ -974,39 +974,23 @@ "source": [ "## 7. Save the adapter and use it at query time\n", "\n", - "The adapter is a single small `.pt` file. At inference you apply it to the\n", + "The adapter is a single small `.safetensors` file. At inference you apply it to the\n", "query embedding **before** your usual nearest-neighbour search — the corpus\n", "index stays untouched." ] }, { "cell_type": "code", - "execution_count": 28, + "execution_count": null, "id": "73eec0e0", - "metadata": { - "execution": { - "iopub.execute_input": "2026-06-15T09:18:35.592659Z", - "iopub.status.busy": "2026-06-15T09:18:35.592594Z", - "iopub.status.idle": "2026-06-15T09:18:35.603505Z", - "shell.execute_reply": "2026-06-15T09:18:35.603170Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Saved adapter -> examples/artifacts/santander_adapter.pt\n", - "query vector shape: (1, 1536) -> adapted: (1, 1536)\n" - ] - } - ], + "metadata": {}, + "outputs": [], "source": [ "from linear_adapter_trainer import LinearAdapter\n", "\n", "out_dir = REPO_ROOT / \"examples\" / \"artifacts\"\n", "out_dir.mkdir(parents=True, exist_ok=True)\n", - "adapter_path = out_dir / \"santander_adapter.pt\"\n", + "adapter_path = out_dir / \"santander_adapter.safetensors\"\n", "result.adapter.save(adapter_path)\n", "print(\"Saved adapter ->\", adapter_path)\n", "\n", diff --git a/linear_adapter_trainer/adapter/model.py b/linear_adapter_trainer/adapter/model.py index c24c4c3..3389630 100644 --- a/linear_adapter_trainer/adapter/model.py +++ b/linear_adapter_trainer/adapter/model.py @@ -12,14 +12,28 @@ from __future__ import annotations +import json +import os +import tempfile +import warnings +from collections.abc import Mapping from dataclasses import asdict, dataclass from pathlib import Path from typing import Any import numpy as np import torch +from safetensors import safe_open +from safetensors.torch import save_file from torch import nn +_FORMAT_VERSION = "1" +_CONFIG_METADATA_KEY = "linear_adapter_config" +_VERSION_METADATA_KEY = "linear_adapter_format_version" +_METADATA_KEYS = {_CONFIG_METADATA_KEY, _VERSION_METADATA_KEY} +_CONFIG_KEYS = {"input_dim", "output_dim", "bias", "residual", "normalize_output"} +_LEGACY_SUFFIXES = {".pt", ".pth"} + @dataclass(slots=True) class AdapterConfig: @@ -90,28 +104,200 @@ def transform(self, embeddings: np.ndarray, *, batch_size: int = 1024) -> np.nda # -- persistence ------------------------------------------------------- def save(self, path: str | Path) -> None: - """Save weights and architecture to a single ``.pt`` file.""" + """Save weights and architecture to a pickle-free safetensors file.""" path = Path(path) + if path.suffix.lower() != ".safetensors": + raise ValueError( + "New checkpoints must use the '.safetensors' extension. " + "Use LinearAdapter.migrate_checkpoint() for legacy '.pt' checkpoints." + ) path.parent.mkdir(parents=True, exist_ok=True) - torch.save({"config": asdict(self.config), "state_dict": self.state_dict()}, path) + + metadata = { + _VERSION_METADATA_KEY: _FORMAT_VERSION, + _CONFIG_METADATA_KEY: json.dumps( + asdict(self.config), sort_keys=True, separators=(",", ":"), allow_nan=False + ), + } + tensors = { + name: tensor.detach().cpu().contiguous() for name, tensor in self.state_dict().items() + } + + file_descriptor, temporary_name = tempfile.mkstemp( + dir=path.parent, prefix=".linear-adapter-", suffix=".tmp" + ) + os.close(file_descriptor) + temporary_path = Path(temporary_name) + try: + save_file(tensors, str(temporary_path), metadata=metadata) + with temporary_path.open("r+b") as handle: + os.fsync(handle.fileno()) + os.replace(temporary_path, path) + if os.name != "nt": + directory_descriptor = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_descriptor) + finally: + os.close(directory_descriptor) + finally: + temporary_path.unlink(missing_ok=True) @classmethod def load(cls, path: str | Path, *, map_location: Any = "cpu") -> LinearAdapter: - """Load an adapter previously written by :meth:`save`. - - Checkpoints are read with ``weights_only=True`` so that ``torch.load`` - never unpickles arbitrary Python objects. The payload written by - :meth:`save` is plain primitives plus tensors, which load correctly - under the safe allowlist. Combined with the ``torch>=2.10.0`` floor in - ``pyproject.toml``, this closes the arbitrary-code-execution vector of - the previous ``weights_only=False`` behavior (CWE-502). - - Security note: only load checkpoints from sources you trust. Treat a - ``.pt`` file like any other executable artifact and verify its origin - (e.g. a SHA-256 checksum) before loading. + """Load a safetensors checkpoint or an explicitly named legacy checkpoint. + + ``.safetensors`` files use strict, versioned JSON metadata and never + invoke pickle. Legacy ``.pt`` and ``.pth`` files remain readable through + PyTorch's restricted ``weights_only=True`` loader and emit a deprecation + warning. Parser selection is extension-based: a malformed safetensors + file is never retried through the legacy loader. ``map_location`` controls + tensor deserialization; as in the legacy implementation, the returned + adapter uses the module's default device and dtype. + + Only load checkpoints from sources you trust and verify artifact + checksums before loading, especially for legacy PyTorch checkpoints. """ + path = Path(path) + suffix = path.suffix.lower() + if suffix == ".safetensors": + return cls._load_safetensors(path, map_location=map_location) + if suffix in _LEGACY_SUFFIXES: + return cls._load_legacy(path, map_location=map_location, warning_stacklevel=3) + raise ValueError( + "Unsupported checkpoint extension. Expected '.safetensors', '.pt', or '.pth'." + ) + + @classmethod + def migrate_checkpoint( + cls, + source: str | Path, + destination: str | Path, + *, + map_location: Any = "cpu", + ) -> LinearAdapter: + """Migrate a legacy ``.pt`` or ``.pth`` checkpoint to safetensors.""" + source = Path(source) + destination = Path(destination) + if source.suffix.lower() not in _LEGACY_SUFFIXES: + raise ValueError("Migration source must use the '.pt' or '.pth' extension.") + if destination.suffix.lower() != ".safetensors": + raise ValueError("Migration destination must use the '.safetensors' extension.") + adapter = cls._load_legacy(source, map_location=map_location, warning_stacklevel=3) + adapter.save(destination) + return adapter + + @classmethod + def _load_safetensors(cls, path: Path, *, map_location: Any) -> LinearAdapter: + if not isinstance(map_location, (str, torch.device)): + raise TypeError("Safetensors map_location must be a string or torch.device.") + device = str(map_location) + with safe_open(str(path), framework="pt", device=device) as checkpoint: + metadata = checkpoint.metadata() + config = _config_from_metadata(metadata) + adapter = cls(config) + expected_keys = set(adapter.state_dict()) + actual_keys = set(checkpoint.keys()) + if actual_keys != expected_keys: + missing = sorted(expected_keys - actual_keys) + unexpected = sorted(actual_keys - expected_keys) + raise ValueError( + f"Invalid checkpoint tensor keys: missing={missing}, unexpected={unexpected}." + ) + state_dict = {name: checkpoint.get_tensor(name) for name in sorted(actual_keys)} + + try: + adapter.load_state_dict(state_dict, strict=True) + except RuntimeError as error: + raise ValueError("Checkpoint tensors do not match the adapter architecture.") from error + adapter.eval() + return adapter + + @classmethod + def _load_legacy( + cls, path: Path, *, map_location: Any, warning_stacklevel: int + ) -> LinearAdapter: + warnings.warn( + "Legacy PyTorch checkpoints are deprecated. Migrate with " + "LinearAdapter.migrate_checkpoint(source, destination).", + DeprecationWarning, + stacklevel=warning_stacklevel, + ) payload = torch.load(path, map_location=map_location, weights_only=True) - adapter = cls(AdapterConfig(**payload["config"])) - adapter.load_state_dict(payload["state_dict"]) + if not isinstance(payload, Mapping) or set(payload) != {"config", "state_dict"}: + raise ValueError("Legacy checkpoint must contain exactly 'config' and 'state_dict'.") + config = _validate_config(payload["config"]) + state_dict = payload["state_dict"] + if not isinstance(state_dict, Mapping) or not all( + isinstance(key, str) and isinstance(value, torch.Tensor) + for key, value in state_dict.items() + ): + raise ValueError("Legacy checkpoint state_dict must map strings to tensors.") + + adapter = cls(config) + expected_keys = set(adapter.state_dict()) + actual_keys = set(state_dict) + if actual_keys != expected_keys: + missing = sorted(expected_keys - actual_keys) + unexpected = sorted(actual_keys - expected_keys) + raise ValueError( + f"Invalid legacy checkpoint tensor keys: missing={missing}, " + f"unexpected={unexpected}." + ) + try: + adapter.load_state_dict(dict(state_dict), strict=True) + except RuntimeError as error: + raise ValueError( + "Legacy checkpoint tensors do not match the adapter architecture." + ) from error adapter.eval() return adapter + + +def _config_from_metadata(metadata: dict[str, str] | None) -> AdapterConfig: + if metadata is None or set(metadata) != _METADATA_KEYS: + actual_keys = sorted(metadata or {}) + raise ValueError( + f"Invalid checkpoint metadata keys: expected={sorted(_METADATA_KEYS)}, " + f"actual={actual_keys}." + ) + version = metadata[_VERSION_METADATA_KEY] + if version != _FORMAT_VERSION: + raise ValueError( + f"Unsupported linear adapter checkpoint format version {version!r}; " + f"supported version is {_FORMAT_VERSION}." + ) + try: + config = json.loads(metadata[_CONFIG_METADATA_KEY]) + except json.JSONDecodeError as error: + raise ValueError("Checkpoint adapter configuration is not valid JSON.") from error + return _validate_config(config) + + +def _validate_config(value: Any) -> AdapterConfig: + if not isinstance(value, Mapping) or set(value) != _CONFIG_KEYS: + actual_keys = sorted(value) if isinstance(value, Mapping) else [] + raise ValueError( + f"Invalid adapter configuration keys: expected={sorted(_CONFIG_KEYS)}, " + f"actual={actual_keys}." + ) + + input_dim = value["input_dim"] + output_dim = value["output_dim"] + if type(input_dim) is not int or input_dim <= 0: + raise ValueError("Adapter configuration input_dim must be a positive integer.") + if output_dim is not None and (type(output_dim) is not int or output_dim <= 0): + raise ValueError("Adapter configuration output_dim must be null or a positive integer.") + for key in ("bias", "residual", "normalize_output"): + if type(value[key]) is not bool: + raise ValueError(f"Adapter configuration {key} must be a boolean.") + + config = AdapterConfig( + input_dim=input_dim, + output_dim=output_dim, + bias=value["bias"], + residual=value["residual"], + normalize_output=value["normalize_output"], + ) + if config.residual and config.resolved_output_dim() != config.input_dim: + raise ValueError("residual adapters require output_dim == input_dim.") + return config diff --git a/linear_adapter_trainer/cli.py b/linear_adapter_trainer/cli.py index 034ec49..d1034be 100644 --- a/linear_adapter_trainer/cli.py +++ b/linear_adapter_trainer/cli.py @@ -40,11 +40,20 @@ def _output(cfg: dict[str, Any]) -> dict[str, Any]: out = cfg.get("output", {}) return { "dataset_dir": out.get("dataset_dir", "artifacts/dataset"), - "adapter_path": out.get("adapter_path", "artifacts/adapter.pt"), + "adapter_path": out.get("adapter_path", "artifacts/adapter.safetensors"), "metrics_path": out.get("metrics_path", "artifacts/metrics.json"), } +def _validate_adapter_output_path(path: str | Path) -> None: + if Path(path).suffix.lower() != ".safetensors": + raise ValueError( + "Change output.adapter_path to use the '.safetensors' extension. " + "Existing '.pt' or '.pth' checkpoints can be converted with " + "LinearAdapter.migrate_checkpoint()." + ) + + def cmd_generate(cfg: dict[str, Any]) -> TripletDataset: kb = build_knowledge_base(cfg.get("knowledge_base", {})) embedder = build_embedder(cfg.get("embedder", {})) @@ -72,6 +81,7 @@ def cmd_generate(cfg: dict[str, Any]) -> TripletDataset: def cmd_train(cfg: dict[str, Any]) -> None: out = _output(cfg) + _validate_adapter_output_path(out["adapter_path"]) kb = build_knowledge_base(cfg.get("knowledge_base", {})) embedder = build_embedder(cfg.get("embedder", {})) dataset = TripletDataset.load(out["dataset_dir"]) @@ -107,6 +117,7 @@ def cmd_evaluate(cfg: dict[str, Any]) -> None: def cmd_run(cfg: dict[str, Any]) -> None: + _validate_adapter_output_path(_output(cfg)["adapter_path"]) cmd_generate(cfg) cmd_train(cfg) diff --git a/pyproject.toml b/pyproject.toml index 20bda5d..8eff3c3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ classifiers = [ ] dependencies = [ "numpy>=2.0.0", + "safetensors>=0.6.2", # >=2.10.0 required: earlier releases have RCE-class flaws reachable via # torch.load even with weights_only=True (CVE-2025-32434, CVE-2026-24747). "torch>=2.10.0", diff --git a/tests/test_adapter_model.py b/tests/test_adapter_model.py index a77073e..69cd9c0 100644 --- a/tests/test_adapter_model.py +++ b/tests/test_adapter_model.py @@ -1,17 +1,29 @@ # Copyright (c) 2026 Santander Group # SPDX-License-Identifier: Apache-2.0 +import json import os import pickle +from dataclasses import asdict +from pathlib import Path import numpy as np import pytest import torch +from safetensors import SafetensorError, safe_open +from safetensors.torch import save_file from linear_adapter_trainer.adapter.losses import TripletLoss from linear_adapter_trainer.adapter.model import AdapterConfig, LinearAdapter +def _metadata(config: AdapterConfig, *, version: str = "1") -> dict[str, str]: + return { + "linear_adapter_format_version": version, + "linear_adapter_config": json.dumps(asdict(config), sort_keys=True, separators=(",", ":")), + } + + def test_residual_adapter_starts_as_identity(): adapter = LinearAdapter(AdapterConfig(input_dim=8, residual=True, normalize_output=False)) x = torch.randn(4, 8) @@ -34,16 +46,350 @@ def test_transform_numpy_shape(): assert out.shape == (5, 6) -def test_save_and_load(tmp_path): - adapter = LinearAdapter(AdapterConfig(input_dim=4, residual=True)) +def test_residual_adapter_rejects_mismatched_output_dimension(): + with pytest.raises(ValueError, match="output_dim == input_dim"): + LinearAdapter(AdapterConfig(input_dim=4, output_dim=3, residual=True)) + + +@pytest.mark.parametrize( + "config", + [ + AdapterConfig(input_dim=4), + AdapterConfig(input_dim=4, bias=False, normalize_output=False), + AdapterConfig(input_dim=4, residual=False, normalize_output=False), + AdapterConfig(input_dim=4, output_dim=3, residual=False), + ], +) +def test_save_and_load(tmp_path, config): + adapter = LinearAdapter(config) with torch.no_grad(): adapter.linear.weight.add_(0.05) - path = tmp_path / "adapter.pt" + path = tmp_path / "adapter.safetensors" adapter.save(path) loaded = LinearAdapter.load(path) x = torch.randn(3, 4) assert torch.allclose(adapter(x), loaded(x), atol=1e-6) + assert loaded.config == config + assert not loaded.training + + with safe_open(path, framework="pt", device="cpu") as checkpoint: + assert checkpoint.metadata() == _metadata(config) + assert set(checkpoint.keys()) == set(adapter.state_dict()) + + +def test_save_rejects_legacy_extension(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + with pytest.raises(ValueError, match="must use the '.safetensors' extension"): + adapter.save(tmp_path / "adapter.pt") + + +def test_save_cleans_partial_temporary_file_on_failure(tmp_path, monkeypatch): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.safetensors" + path.write_bytes(b"existing") + + def fail_save(tensors, filename, *, metadata): + del tensors, metadata + Path(filename).write_bytes(b"partial") + raise RuntimeError("simulated write failure") + + monkeypatch.setattr("linear_adapter_trainer.adapter.model.save_file", fail_save) + with pytest.raises(RuntimeError, match="simulated write failure"): + adapter.save(path) + assert path.read_bytes() == b"existing" + assert list(tmp_path.iterdir()) == [path] + + +def test_save_atomically_replaces_existing_checkpoint(tmp_path, monkeypatch): + path = tmp_path / "adapter.safetensors" + original = LinearAdapter(AdapterConfig(input_dim=4, normalize_output=False)) + original.save(path) + original_bytes = path.read_bytes() + replace_observations = [] + real_replace = os.replace + + def observe_replace(source, destination): + replace_observations.append( + { + "source_is_temporary": source != destination and source.exists(), + "destination_is_unchanged": destination.read_bytes() == original_bytes, + } + ) + real_replace(source, destination) + + monkeypatch.setattr("linear_adapter_trainer.adapter.model.os.replace", observe_replace) + + replacement = LinearAdapter(AdapterConfig(input_dim=4, normalize_output=False)) + with torch.no_grad(): + replacement.linear.weight.add_(0.25) + replacement.save(path) + + loaded = LinearAdapter.load(path) + assert replace_observations == [{"source_is_temporary": True, "destination_is_unchanged": True}] + for key, tensor in replacement.state_dict().items(): + assert torch.equal(tensor, loaded.state_dict()[key]) + + +def test_load_rejects_unsupported_format_version(tmp_path): + config = AdapterConfig(input_dim=4) + adapter = LinearAdapter(config) + path = tmp_path / "adapter.safetensors" + save_file(adapter.state_dict(), path, metadata=_metadata(config, version="2")) + + with pytest.raises(ValueError, match="Unsupported.*format version '2'"): + LinearAdapter.load(path) + + +@pytest.mark.parametrize( + ("metadata", "message"), + [ + ({}, "Invalid checkpoint metadata keys"), + ({"linear_adapter_format_version": "1"}, "Invalid checkpoint metadata keys"), + ( + { + **_metadata(AdapterConfig(input_dim=4)), + "unexpected": "metadata", + }, + "Invalid checkpoint metadata keys", + ), + ( + { + "linear_adapter_format_version": "1", + "linear_adapter_config": "not-json", + }, + "not valid JSON", + ), + ], +) +def test_load_rejects_invalid_metadata(tmp_path, metadata, message): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.safetensors" + save_file(adapter.state_dict(), path, metadata=metadata) + + with pytest.raises(ValueError, match=message): + LinearAdapter.load(path) + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda config: config.pop("bias"), "Invalid adapter configuration keys"), + (lambda config: config.update(extra=True), "Invalid adapter configuration keys"), + (lambda config: config.update(input_dim=True), "input_dim must be a positive integer"), + (lambda config: config.update(input_dim=-1), "input_dim must be a positive integer"), + (lambda config: config.update(input_dim=4.0), "input_dim must be a positive integer"), + (lambda config: config.update(input_dim="4"), "input_dim must be a positive integer"), + ( + lambda config: config.update(output_dim=0), + "output_dim must be null or a positive integer", + ), + ( + lambda config: config.update(output_dim=-1), + "output_dim must be null or a positive integer", + ), + ( + lambda config: config.update(output_dim=4.0), + "output_dim must be null or a positive integer", + ), + ( + lambda config: config.update(output_dim="4"), + "output_dim must be null or a positive integer", + ), + (lambda config: config.update(residual=1), "residual must be a boolean"), + ( + lambda config: config.update(output_dim=3, residual=True), + "residual adapters require output_dim == input_dim", + ), + ], +) +def test_load_rejects_invalid_config(tmp_path, mutate, message): + config = asdict(AdapterConfig(input_dim=4)) + mutate(config) + metadata = _metadata(AdapterConfig(input_dim=4)) + metadata["linear_adapter_config"] = json.dumps(config) + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.safetensors" + save_file(adapter.state_dict(), path, metadata=metadata) + + with pytest.raises(ValueError, match=message): + LinearAdapter.load(path) + + +def test_load_rejects_missing_unexpected_and_mismatched_tensors(tmp_path): + config = AdapterConfig(input_dim=4) + metadata = _metadata(config) + + missing_path = tmp_path / "missing.safetensors" + save_file({"linear.weight": torch.zeros(4, 4)}, missing_path, metadata=metadata) + with pytest.raises(ValueError, match=r"missing=\['linear.bias'\]"): + LinearAdapter.load(missing_path) + + unexpected_path = tmp_path / "unexpected.safetensors" + state_dict = LinearAdapter(config).state_dict() + state_dict["extra"] = torch.zeros(1) + save_file(state_dict, unexpected_path, metadata=metadata) + with pytest.raises(ValueError, match=r"unexpected=\['extra'\]"): + LinearAdapter.load(unexpected_path) + + mismatched_path = tmp_path / "mismatched.safetensors" + save_file( + {"linear.weight": torch.zeros(3, 4), "linear.bias": torch.zeros(4)}, + mismatched_path, + metadata=metadata, + ) + with pytest.raises(ValueError, match="tensors do not match"): + LinearAdapter.load(mismatched_path) + + +def test_load_rejects_corrupt_safetensors_without_legacy_fallback(tmp_path): + marker = tmp_path / "pwned-via-fallback.txt" + + class _Evil: + def __reduce__(self): + return (os.system, (f"touch {marker}",)) + + path = tmp_path / "evil.safetensors" + torch.save(_Evil(), path) + + with pytest.raises(SafetensorError): + LinearAdapter.load(path) + assert not marker.exists() + + +def test_load_rejects_unknown_extension(tmp_path): + with pytest.raises(ValueError, match="Unsupported checkpoint extension"): + LinearAdapter.load(tmp_path / "adapter.bin") + + +def test_safetensors_rejects_unsupported_map_location(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.safetensors" + adapter.save(path) + + with pytest.raises(TypeError, match="string or torch.device"): + LinearAdapter.load(path, map_location={"cpu": "cpu"}) + + +def test_safetensors_accepts_torch_device_map_location(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.safetensors" + adapter.save(path) + + loaded = LinearAdapter.load(path, map_location=torch.device("cpu")) + assert loaded.config == adapter.config + + +def test_safetensors_accepts_string_map_location_and_preserves_legacy_dtype_semantics(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)).double() + path = tmp_path / "adapter.safetensors" + adapter.save(str(path)) + + loaded = LinearAdapter.load(str(path), map_location="cpu") + assert next(loaded.parameters()).device.type == "cpu" + assert next(loaded.parameters()).dtype == torch.get_default_dtype() + + +def test_safetensors_extension_is_case_insensitive(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.SAFETENSORS" + adapter.save(path) + + loaded = LinearAdapter.load(path) + assert loaded.config == adapter.config + + +def test_migration_rejects_non_legacy_source(tmp_path): + with pytest.raises(ValueError, match="Migration source"): + LinearAdapter.migrate_checkpoint( + tmp_path / "adapter.safetensors", tmp_path / "migrated.safetensors" + ) + + +def test_migration_rejects_legacy_destination_extension(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + source = tmp_path / "adapter.pt" + torch.save({"config": asdict(adapter.config), "state_dict": adapter.state_dict()}, source) + + with pytest.raises(ValueError, match="Migration destination"): + LinearAdapter.migrate_checkpoint(source, tmp_path / "migrated.pt") + + +def test_loads_and_migrates_legacy_checkpoint(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4, normalize_output=False)) + with torch.no_grad(): + adapter.linear.weight.add_(0.05) + legacy_path = tmp_path / "adapter.pt" + torch.save({"config": asdict(adapter.config), "state_dict": adapter.state_dict()}, legacy_path) + + with pytest.warns( + DeprecationWarning, match="Legacy PyTorch checkpoints are deprecated" + ) as load_warnings: + loaded = LinearAdapter.load(legacy_path) + assert load_warnings[0].filename == __file__ + assert loaded.config == adapter.config + for key, tensor in adapter.state_dict().items(): + assert torch.equal(tensor, loaded.state_dict()[key]) + + destination = tmp_path / "adapter.safetensors" + with pytest.warns(DeprecationWarning) as migration_warnings: + migrated = LinearAdapter.migrate_checkpoint(legacy_path, destination) + assert migration_warnings[0].filename == __file__ + reloaded = LinearAdapter.load(destination) + assert migrated.config == reloaded.config == adapter.config + for key, tensor in adapter.state_dict().items(): + assert torch.equal(tensor, reloaded.state_dict()[key]) + + +def test_load_rejects_invalid_legacy_payload(tmp_path): + adapter = LinearAdapter(AdapterConfig(input_dim=4)) + path = tmp_path / "adapter.pt" + torch.save( + { + "config": asdict(adapter.config), + "state_dict": adapter.state_dict(), + "unexpected": True, + }, + path, + ) + + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match="exactly"): + LinearAdapter.load(path) + + +def test_load_rejects_invalid_legacy_state_dict(tmp_path): + config = AdapterConfig(input_dim=4) + invalid_value_path = tmp_path / "invalid-value.pt" + torch.save( + {"config": asdict(config), "state_dict": {"linear.weight": "no"}}, invalid_value_path + ) + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match="strings to tensors"): + LinearAdapter.load(invalid_value_path) + + missing_key_path = tmp_path / "missing-key.pt" + torch.save( + {"config": asdict(config), "state_dict": {"linear.weight": torch.zeros(4, 4)}}, + missing_key_path, + ) + with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match="missing=.*linear.bias"): + LinearAdapter.load(missing_key_path) + + mismatched_path = tmp_path / "mismatched.pt" + torch.save( + { + "config": asdict(config), + "state_dict": { + "linear.weight": torch.zeros(3, 4), + "linear.bias": torch.zeros(4), + }, + }, + mismatched_path, + ) + with ( + pytest.warns(DeprecationWarning), + pytest.raises(ValueError, match="Legacy checkpoint tensors do not match"), + ): + LinearAdapter.load(mismatched_path) def test_load_rejects_malicious_pickle(tmp_path): @@ -61,7 +407,7 @@ def __reduce__(self): evil_path = tmp_path / "evil.pt" torch.save({"config": {"input_dim": 4}, "state_dict": {}, "x": _Evil()}, evil_path) - with pytest.raises(pickle.UnpicklingError): + with pytest.warns(DeprecationWarning), pytest.raises(pickle.UnpicklingError): LinearAdapter.load(evil_path) assert not marker.exists() diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..769e138 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 Santander Group +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from linear_adapter_trainer.cli import _output, build_parser, cmd_run, cmd_train + + +def test_output_defaults_to_safetensors(): + assert _output({})["adapter_path"] == "artifacts/adapter.safetensors" + + +@pytest.mark.parametrize("command", ["generate", "train", "evaluate", "run"]) +def test_parser_accepts_config_for_each_command(command): + args = build_parser().parse_args([command, "config.toml"]) + assert args.command == command + assert args.config == "config.toml" + + +@pytest.mark.parametrize("command", [cmd_train, cmd_run]) +def test_training_commands_reject_legacy_output_before_work(command, monkeypatch): + started = False + + def fail_if_started(*args, **kwargs): + nonlocal started + started = True + raise AssertionError("training work started before output validation") + + monkeypatch.setattr("linear_adapter_trainer.cli.build_knowledge_base", fail_if_started) + with pytest.raises(ValueError, match=r"Change output\.adapter_path"): + command({"output": {"adapter_path": "artifacts/adapter.pt"}}) + assert not started