wehub-resource-sync
54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any, BinaryIO, Literal, overload
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
from .._serializable import Deserializer, Serializable, Serializer
|
|
from ..utils import safe_isinstance
|
|
|
|
|
|
class Model(Serializable):
|
|
"""This is the superclass of all models."""
|
|
|
|
def __init__(self, model: Any = None) -> None:
|
|
"""Wrap a callable model as a SHAP Model object."""
|
|
if isinstance(model, Model):
|
|
self.inner_model: Any = model.inner_model
|
|
else:
|
|
self.inner_model = model
|
|
|
|
if hasattr(model, "output_names"):
|
|
self.output_names = model.output_names
|
|
|
|
def __call__(self, *args: Any) -> npt.NDArray[Any]:
|
|
out = self.inner_model(*args)
|
|
is_tensor = safe_isinstance(out, "torch.Tensor")
|
|
out = out.cpu().detach().numpy() if is_tensor else np.array(out)
|
|
return out
|
|
|
|
def save(self, out_file: BinaryIO) -> None:
|
|
"""Save the model to the given file stream."""
|
|
super().save(out_file)
|
|
with Serializer(out_file, "shap.Model", version=0) as s:
|
|
s.save("model", self.inner_model)
|
|
|
|
@overload
|
|
@classmethod
|
|
def load(cls, in_file: BinaryIO, instantiate: Literal[True] = True) -> Model: ...
|
|
|
|
@overload
|
|
@classmethod
|
|
def load(cls, in_file: BinaryIO, instantiate: Literal[False]) -> dict[str, Any]: ...
|
|
|
|
@classmethod
|
|
def load(cls, in_file: BinaryIO, instantiate: bool = True) -> Model | dict[str, Any]:
|
|
if instantiate:
|
|
return cls._instantiated_load(in_file)
|
|
|
|
kwargs = super().load(in_file, instantiate=False)
|
|
with Deserializer(in_file, "shap.Model", min_version=0, max_version=0) as s:
|
|
kwargs["model"] = s.load("model")
|
|
return kwargs # type: ignore[return-value]
|