| """Explicitly trusted legacy checkpoint loading and guarded inference. |
| |
| Pickle executes code. Namespace compatibility and SHA256 comparison do not make |
| an untrusted pickle safe; users must opt in before any unpickling takes place. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import hmac |
| import pickle |
| import re |
| from pathlib import Path |
| from typing import Any, Mapping |
|
|
|
|
| class WISPUnpickler(pickle.Unpickler): |
| """Map historical method modules, without installing process-wide aliases. |
| |
| This is a compatibility loader, NOT a restricted or sandboxed unpickler. |
| """ |
|
|
| def find_class(self, module: str, name: str) -> Any: |
| if module == "tsevolve.search.ts_evolve_v3": |
| module = "wisp.search.controller" |
| elif module == "tsevolve" or module.startswith("tsevolve."): |
| module = "wisp" + module[len("tsevolve"):] |
| return super().find_class(module, name) |
|
|
|
|
| def sha256_file(path: str | Path) -> str: |
| digest = hashlib.sha256() |
| with Path(path).open("rb") as stream: |
| for block in iter(lambda: stream.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def load_checkpoint(path: str | Path, *, trust_pickle: bool = False, |
| expected_sha256: str | None = None) -> Any: |
| """Load a trusted fitted WISP estimator, optionally checking its SHA256. |
| |
| Supply ``expected_sha256`` from an independently trusted source. A digest |
| supplied alongside a checkpoint proves integrity, not publisher identity. |
| Historical checkpoint files are left unchanged on disk. |
| """ |
| if not trust_pickle: |
| raise ValueError("Pickle can execute arbitrary code. Load only trusted WISP " |
| "checkpoints and explicitly set trust_pickle=True.") |
| path = Path(path) |
| if expected_sha256 is not None: |
| if not re.fullmatch(r"[0-9a-fA-F]{64}", expected_sha256): |
| raise ValueError("expected_sha256 must be a 64-character hexadecimal digest") |
| if not hmac.compare_digest(sha256_file(path), expected_sha256.lower()): |
| raise ValueError(f"SHA256 mismatch for {path.name}; refusing to unpickle") |
| with path.open("rb") as stream: |
| model = WISPUnpickler(stream).load() |
| if not callable(getattr(model, "predict", None)): |
| raise TypeError("Checkpoint does not contain a fitted estimator with predict()") |
| return model |
|
|
|
|
| def _array_vector(value: Any, *, name: str, n: int) -> Any: |
| import numpy as np |
|
|
| array = np.asarray(value) |
| if array.ndim != 1 or len(array) != n: |
| raise ValueError(f"{name} must be a one-dimensional vector with {n} entries") |
| return array |
|
|
|
|
| def _check_X(X: Any, metadata: Mapping[str, Any] | None = None) -> Any: |
| import numpy as np |
|
|
| array = np.asarray(X, dtype=np.float64) |
| if array.ndim != 3 or not all(array.shape): |
| raise ValueError("X must have nonempty shape (n_windows, n_timepoints, n_channels)") |
| if not np.isfinite(array).all(): |
| raise ValueError("X contains NaN or infinite values; apply the dataset preprocessing first") |
| if metadata and metadata.get("input_shape") is not None: |
| expected = tuple(int(v) for v in metadata["input_shape"]) |
| |
| if len(expected) == 3: |
| expected = expected[1:] |
| if len(expected) != 2 or array.shape[1:] != expected: |
| raise ValueError(f"Checkpoint expects time/channel shape {expected}; got {array.shape[1:]}") |
| return array |
|
|
|
|
| def _validate_class_axis(model: Any, metadata: Mapping[str, Any] | None) -> None: |
| """Never redefine a fitted global class axis using a shorter input vocabulary.""" |
| import numpy as np |
|
|
| classes = getattr(model, "classes_", None) |
| if classes is not None: |
| classes = np.asarray(classes) |
| if (classes.ndim != 1 or not len(classes) |
| or not np.issubdtype(classes.dtype, np.integer) |
| or not np.array_equal(classes, np.arange(len(classes)))): |
| raise ValueError("Checkpoint classes_ must be the contiguous encoded global " |
| "axis 0..K-1; refusing to silently remap a nonstandard class axis") |
| names = None if metadata is None else metadata.get("label_names") |
| if names is not None: |
| if isinstance(names, (str, bytes)) or np.asarray(names).ndim != 1: |
| raise ValueError("label_names must be a one-dimensional global class vocabulary") |
| if classes is not None and len(names) != len(classes): |
| raise ValueError(f"label_names has {len(names)} entries but checkpoint global " |
| f"classes_ has {len(classes)}; preserve the complete fitted vocabulary") |
|
|
|
|
| def require_sequence_metadata(model: Any, subject: Any, time_index: Any, *, n: int) -> tuple[Any, Any]: |
| """Do not silently bypass the sequence decoder stored in a checkpoint.""" |
| needs_sequence = getattr(model, "smoother_", None) is not None |
| if needs_sequence and (subject is None or time_index is None): |
| raise ValueError("This checkpoint uses the sequence decoder: both subject " |
| "(sequence/session IDs) and time_index are required. " |
| "Do not merge unrelated recording sessions into one sequence.") |
| if subject is not None: |
| subject = _array_vector(subject, name="subject", n=n) |
| if time_index is not None: |
| time_index = _array_vector(time_index, name="time_index", n=n) |
| return subject, time_index |
|
|
|
|
| def set_inference_jobs(model: Any, n_jobs: int) -> None: |
| """Limit the known fitted estimator/ensemble members without refitting.""" |
| if n_jobs < 1: |
| raise ValueError("n_jobs must be positive") |
| seen: set[int] = set() |
|
|
| def visit(estimator: Any) -> None: |
| if estimator is None or id(estimator) in seen: |
| return |
| seen.add(id(estimator)) |
| if hasattr(estimator, "n_jobs"): |
| estimator.n_jobs = n_jobs |
| for attr in ("model_", "members_", "estimators_", "steps"): |
| child = getattr(estimator, attr, None) |
| if isinstance(child, (tuple, list)): |
| for item in child: |
| visit(item[-1] if isinstance(item, tuple) else item) |
| elif child is not None: |
| visit(child) |
|
|
| visit(model) |
|
|
|
|
| def predict_checkpoint(model: Any, X: Any, *, subject: Any = None, |
| time_index: Any = None, metadata: Mapping[str, Any] | None = None, |
| n_jobs: int = 1) -> Any: |
| """Predict encoded class IDs using the fitted checkpoint's actual decoder. |
| |
| ``metadata`` may contain ``input_shape=[T,C]`` and ``label_names``; it never |
| authorizes pickle loading, downloads data, or silently changes a decoder. |
| """ |
| import numpy as np |
|
|
| _validate_class_axis(model, metadata) |
| X = _check_X(X, metadata) |
| subject, time_index = require_sequence_metadata(model, subject, time_index, n=len(X)) |
| set_inference_jobs(model, n_jobs) |
| prediction = np.asarray(model.predict(X, groups=subject, time_index=time_index)) |
| if prediction.ndim != 1 or len(prediction) != len(X): |
| raise ValueError("Model returned an invalid prediction shape") |
| if not np.issubdtype(prediction.dtype, np.integer): |
| raise ValueError("WISP checkpoints must return encoded integer class IDs") |
| prediction = prediction.astype(np.int64, copy=False) |
| names = None if metadata is None else metadata.get("label_names") |
| if names is not None and len(prediction) and (prediction.min() < 0 or prediction.max() >= len(names)): |
| raise ValueError("Predictions are incompatible with the supplied label_names") |
| return prediction |
|
|
