wehub-resource-sync
202 lines
8.6 KiB
Python
202 lines
8.6 KiB
Python
import inspect
|
|
import logging
|
|
import pickle
|
|
|
|
import cloudpickle
|
|
import numpy as np
|
|
|
|
log = logging.getLogger("shap")
|
|
|
|
|
|
class Serializable:
|
|
"""This is the superclass of all serializable objects."""
|
|
|
|
def save(self, out_file):
|
|
"""Save the model to the given file stream."""
|
|
pickle.dump(type(self), out_file)
|
|
|
|
@classmethod
|
|
def load(cls, in_file, instantiate=True):
|
|
"""This is meant to be overridden by subclasses and called with super.
|
|
|
|
We return constructor argument values when not being instantiated. Since there are no
|
|
constructor arguments for the Serializable class we just return an empty dictionary.
|
|
"""
|
|
if instantiate:
|
|
return cls._instantiated_load(in_file)
|
|
return {}
|
|
|
|
@classmethod
|
|
def _instantiated_load(cls, in_file, **kwargs):
|
|
"""This is meant to be overridden by subclasses and called with super.
|
|
|
|
We return constructor argument values (we have no values to load in this abstract class).
|
|
"""
|
|
obj_type = pickle.load(in_file)
|
|
if obj_type is None:
|
|
return None
|
|
|
|
if not inspect.isclass(obj_type) or (not issubclass(obj_type, cls) and (obj_type is not cls)):
|
|
raise Exception(f"Invalid object type loaded from file. {obj_type} is not a subclass of {cls}.")
|
|
|
|
# here we call the constructor with all the arguments we have loaded
|
|
constructor_args = obj_type.load(in_file, instantiate=False, **kwargs)
|
|
used_args = inspect.getfullargspec(obj_type.__init__)[0]
|
|
return obj_type(**{k: constructor_args[k] for k in constructor_args if k in used_args})
|
|
|
|
|
|
class Serializer:
|
|
"""Save data items to an input stream."""
|
|
|
|
def __init__(self, out_stream, block_name, version):
|
|
self.out_stream = out_stream
|
|
self.block_name = block_name
|
|
self.block_version = version
|
|
self.serializer_version = 0 # update this when the serializer changes
|
|
|
|
def __enter__(self):
|
|
log.debug("serializer_version = %d", self.serializer_version)
|
|
pickle.dump(self.serializer_version, self.out_stream)
|
|
log.debug("block_name = %s", self.block_name)
|
|
pickle.dump(self.block_name, self.out_stream)
|
|
log.debug("block_version = %d", self.block_version)
|
|
pickle.dump(self.block_version, self.out_stream)
|
|
return self
|
|
|
|
def __exit__(self, exception_type, exception_value, traceback):
|
|
log.debug("END_BLOCK___")
|
|
pickle.dump("END_BLOCK___", self.out_stream)
|
|
|
|
def save(self, name, value, encoder="auto"):
|
|
"""Dump a data item to the current input stream."""
|
|
log.debug("name = %s", name)
|
|
pickle.dump(name, self.out_stream)
|
|
if encoder is None or encoder is False:
|
|
log.debug("encoder_name = %s", "no_encoder")
|
|
pickle.dump("no_encoder", self.out_stream)
|
|
elif callable(encoder):
|
|
log.debug("encoder_name = %s", "custom_encoder")
|
|
pickle.dump("custom_encoder", self.out_stream)
|
|
encoder(value, self.out_stream)
|
|
elif encoder == ".save" or (isinstance(value, Serializable) and encoder == "auto"):
|
|
if hasattr(value, "save"):
|
|
# If it can save, use the original serializable.save logic
|
|
log.debug("encoder_name = %s", "serializable.save")
|
|
pickle.dump("serializable.save", self.out_stream)
|
|
value.save(self.out_stream)
|
|
else:
|
|
# If it can't, use the cloudpickle fallback logic directly
|
|
log.debug("encoder_name = %s", "cloudpickle.dump (fallback)")
|
|
pickle.dump("cloudpickle.dump", self.out_stream)
|
|
cloudpickle.dump(value, self.out_stream)
|
|
elif encoder == "auto":
|
|
if isinstance(value, (int, float, str)):
|
|
log.debug("encoder_name = %s", "pickle.dump")
|
|
pickle.dump("pickle.dump", self.out_stream)
|
|
pickle.dump(value, self.out_stream)
|
|
else:
|
|
log.debug("encoder_name = %s", "cloudpickle.dump")
|
|
pickle.dump("cloudpickle.dump", self.out_stream)
|
|
cloudpickle.dump(value, self.out_stream)
|
|
else:
|
|
raise ValueError(f"Unknown encoder type '{encoder}' given for serialization!")
|
|
log.debug("value = %s", str(value))
|
|
|
|
|
|
class Deserializer:
|
|
"""Load data items from an input stream."""
|
|
|
|
def __init__(self, in_stream, block_name, min_version, max_version):
|
|
self.in_stream = in_stream
|
|
self.block_name = block_name
|
|
self.block_min_version = min_version
|
|
self.block_max_version = max_version
|
|
|
|
# update these when the serializer changes
|
|
self.serializer_min_version = 0
|
|
self.serializer_max_version = 0
|
|
|
|
def __enter__(self):
|
|
# confirm the serializer version
|
|
serializer_version = pickle.load(self.in_stream)
|
|
log.debug("serializer_version = %d", serializer_version)
|
|
if serializer_version < self.serializer_min_version:
|
|
raise ValueError(
|
|
f"The file being loaded was saved with a serializer version of {serializer_version}, "
|
|
f"but the current deserializer in SHAP requires at least version {self.serializer_min_version}."
|
|
)
|
|
if serializer_version > self.serializer_max_version:
|
|
raise ValueError(
|
|
f"The file being loaded was saved with a serializer version of {serializer_version}, "
|
|
f"but the current deserializer in SHAP only support up to version {self.serializer_max_version}."
|
|
)
|
|
|
|
# confirm the block name
|
|
block_name = pickle.load(self.in_stream)
|
|
log.debug("block_name = %s", block_name)
|
|
if block_name != self.block_name:
|
|
raise ValueError(
|
|
f"The next data block in the file being loaded was supposed to be {self.block_name}, "
|
|
f"but the next block found was {block_name}."
|
|
)
|
|
|
|
# confirm the block version
|
|
block_version = pickle.load(self.in_stream)
|
|
log.debug("block_version = %d", block_version)
|
|
if block_version < self.block_min_version:
|
|
raise ValueError(
|
|
f"The file being loaded was saved with a block version of {block_version}, "
|
|
f"but the current deserializer in SHAP requires at least version {self.block_min_version}."
|
|
)
|
|
if block_version > self.block_max_version:
|
|
raise ValueError(
|
|
f"The file being loaded was saved with a block version of {block_version}, "
|
|
f"but the current deserializer in SHAP only support up to version {self.block_max_version}."
|
|
)
|
|
return self
|
|
|
|
def __exit__(self, exception_type, exception_value, traceback):
|
|
# confirm the block end token
|
|
for _ in range(100):
|
|
end_token = pickle.load(self.in_stream)
|
|
log.debug("end_token = %s", end_token)
|
|
if end_token == "END_BLOCK___":
|
|
return
|
|
self._load_data_value()
|
|
raise ValueError(f"The data block end token was not found for the block {self.block_name}.")
|
|
|
|
def load(self, name, decoder=None):
|
|
"""Load a data item from the current input stream."""
|
|
# confirm the block name
|
|
loaded_name = pickle.load(self.in_stream)
|
|
log.debug("loaded_name = %s", loaded_name)
|
|
|
|
if loaded_name != name:
|
|
raise ValueError(
|
|
f"The next data item in the file being loaded was supposed to be {name}, "
|
|
f"but the next block found was {loaded_name}."
|
|
) # We should eventually add support for skipping over unused data items in old formats...
|
|
|
|
value = self._load_data_value(decoder)
|
|
log.debug("value = %s", str(value))
|
|
return value
|
|
|
|
def _load_data_value(self, decoder=None):
|
|
encoder_name = pickle.load(self.in_stream)
|
|
log.debug("encoder_name = %s", encoder_name)
|
|
if encoder_name == "custom_encoder" or callable(decoder):
|
|
assert callable(decoder), "You must provide a callable custom decoder for the data item {name}!"
|
|
return decoder(self.in_stream)
|
|
if encoder_name == "no_encoder":
|
|
return None
|
|
if encoder_name == "serializable.save":
|
|
return Serializable.load(self.in_stream)
|
|
if encoder_name == "numpy.save":
|
|
return np.load(self.in_stream)
|
|
if encoder_name == "pickle.dump":
|
|
return pickle.load(self.in_stream)
|
|
if encoder_name == "cloudpickle.dump":
|
|
return cloudpickle.load(self.in_stream)
|
|
|
|
raise ValueError(f"Unsupported encoder type found: {encoder_name}")
|