feat: add pickle-free safetensors checkpoints

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.
 

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__`.
 

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.
 

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`):

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"

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",

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