molcraft 0.1.0rc9__py3-none-any.whl

molcraft/losses.py

@@ -0,0 +1,37 @@
+import warnings
+import keras 
+import numpy as np
+
+
+@keras.saving.register_keras_serializable(package='molcraft')
+class GaussianNegativeLogLikelihood(keras.losses.Loss):
+
+    def __init__(
+        self, 
+        events: int = 1, 
+        name="gaussian_nll", 
+        **kwargs
+    ):
+        super().__init__(name=name, **kwargs)
+        self.events = events
+    
+    def call(self, y_true, y_pred):
+        mean = y_pred[..., :self.events]
+        scale = y_pred[..., self.events:]
+        variance = keras.ops.square(scale)
+        expected_rank = len(keras.ops.shape(mean))
+        current_rank = len(keras.ops.shape(y_true))
+        for _ in range(expected_rank - current_rank):
+            y_true = keras.ops.expand_dims(y_true, axis=-1)
+        return keras.ops.mean(
+            0.5 * keras.ops.log(2.0 * np.pi * variance) + 
+            0.5 * keras.ops.square(y_true - mean) / variance 
+        )
+
+    def get_config(self):
+        config = super().get_config()
+        config['events'] = self.events 
+        return config 
+    
+
+GaussianNLL = GaussianNegativeLogLikelihood

molcraft/models.py

@@ -0,0 +1,623 @@
+import warnings
+import typing
+import keras
+import numpy as np
+import tensorflow as tf
+from pathlib import Path
+from keras.src.models import functional
+
+from molcraft import layers
+from molcraft import tensors
+from molcraft import ops
+
+
+@keras.saving.register_keras_serializable(package="molcraft")
+class GraphModel(layers.GraphLayer, keras.models.Model):
+    
+    """A graph model.
+
+    Currently, the `GraphModel` only supports `GraphTensor` input.
+
+    Build a subclassed GraphModel:
+
+    >>> import molcraft
+    >>> import keras
+    >>>
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
+    >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
+    >>>
+    >>> @keras.saving.register_keras_serializable()
+    >>> class GraphNeuralNetwork(molcraft.models.GraphModel):
+    ...     def __init__(self, units, **kwargs):
+    ...         super().__init__(**kwargs)
+    ...         self.units = units
+    ...         self.node_embedding = molcraft.layers.NodeEmbedding(self.units)
+    ...         self.edge_embedding = molcraft.layers.EdgeEmbedding(self.units)
+    ...         self.conv_1 = molcraft.layers.GraphConv(self.units)
+    ...         self.conv_2 = molcraft.layers.GraphConv(self.units)
+    ...         self.readout = molcraft.layers.Readout('mean')
+    ...         self.dense = keras.layers.Dense(1)
+    ...     def propagate(self, graph):
+    ...         x = self.edge_embedding(self.node_embedding(graph))
+    ...         x = self.conv_2(self.conv_1(x))
+    ...         return self.dense(self.readout(x))
+    ...     def get_config(self):
+    ...         config = super().get_config()
+    ...         config['units'] = self.units
+    ...         return config
+    >>>
+    >>> model = GraphNeuralNetwork(128)
+    >>> model.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-3),
+    ...     loss=keras.losses.MeanSquaredError(),
+    ...     metrics=[keras.metrics.MeanAbsolutePercentageError(name='mape')]
+    ... )
+    >>> model.fit(graph, epochs=10)
+    >>> mse, mape = model.evaluate(graph)
+    >>> preds = model.predict(graph)
+
+    Build a functional GraphModel:
+
+    >>> import molcraft
+    >>> import keras
+    >>>
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
+    >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
+    >>>
+    >>> inputs = molcraft.layers.Input(graph.spec)
+    >>> x = molcraft.layers.NodeEmbedding(128)(inputs)
+    >>> x = molcraft.layers.EdgeEmbedding(128)(x)
+    >>> x = molcraft.layers.GraphConv(128)(x)
+    >>> x = molcraft.layers.GraphConv(128)(x)
+    >>> x = molcraft.layers.Readout('mean')(x)
+    >>> outputs = keras.layers.Dense(1)(x)
+    >>> model = molcraft.models.GraphModel(inputs, outputs)
+    >>> model.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-3),
+    ...     loss=keras.losses.MeanSquaredError(),
+    ...     metrics=[keras.metrics.MeanAbsolutePercentageError(name='mape')]
+    ... )
+    >>> model.fit(graph, epochs=10)
+    >>> mse, mape = model.evaluate(graph)
+    >>> preds = model.predict(graph)
+
+    Build a GraphModel using `from_layers`:
+
+    >>> import molcraft 
+    >>> import keras
+    >>> 
+    >>> featurizer = molcraft.featurizers.MolGraphFeaturizer()
+    >>> graph = featurizer([('N[C@@H](C)C(=O)O', 1.0), ('N[C@@H](CS)C(=O)O', 2.0)])
+    >>> 
+    >>> model = molcraft.models.GraphModel.from_layers([
+    ...     molcraft.layers.Input(graph.spec),
+    ...     molcraft.layers.NodeEmbedding(128),
+    ...     molcraft.layers.EdgeEmbedding(128),
+    ...     molcraft.layers.GraphConv(128),
+    ...     molcraft.layers.GraphConv(128),
+    ...     molcraft.layers.Readout('mean'),
+    ...     keras.layers.Dense(1)
+    ... ])
+    >>> model.compile(
+    ...     optimizer=keras.optimizers.Adam(1e-3),
+    ...     loss=keras.losses.MeanSquaredError(),
+    ...     metrics=[keras.metrics.MeanAbsolutePercentageError(name='mape')]
+    ... )
+    >>> model.fit(graph, epochs=10)
+    >>> mse, mape = model.evaluate(graph)
+    >>> preds = model.predict(graph)
+
+    """
+
+    def __new__(cls, *args, **kwargs):
+        if _functional_init_arguments(args, kwargs) and cls == GraphModel:
+            return FunctionalGraphModel(*args, **kwargs)
+        return super().__new__(cls)
+    
+    def __init__(self, *args, **kwargs):
+        self._model_layers = kwargs.pop('model_layers', None)
+        super().__init__(*args, **kwargs)
+        self.jit_compile = False 
+
+    @classmethod
+    def from_layers(cls, graph_layers: list, **kwargs):
+        """Creates a graph model from a list of graph layers.
+
+        Currently requires `molcraft.layers.Input(spec)`. 
+
+        If `molcraft.layers.Input(spec)` is supplied, it both 
+        creates and builds the layer, as a functional model.
+        `molcraft.layers.Input` is a function which returns
+        a nested structure of graph components based on `spec`.
+
+        Args:
+            graph_layers:
+                A list of `GraphLayer` instances, except the initial element
+                which is a dictionary of Keras tensors produced by 
+                `molcraft.layers.Input(spec)`.
+        """
+        if not tensors.is_graph(graph_layers[0]):
+            return cls(model_layers=graph_layers, **kwargs)
+        elif cls != GraphModel:
+            return cls(model_layers=graph_layers[1:], **kwargs)
+        inputs: dict = graph_layers.pop(0)
+        x = inputs
+        for layer in graph_layers:
+            if isinstance(layer, list):
+                layer = layers.GraphNetwork(layer)
+            x = layer(x)
+        outputs = x
+        return cls(inputs=inputs, outputs=outputs, **kwargs)
+
+    def propagate(self, graph: tensors.GraphTensor) -> tensors.GraphTensor:
+        if self._model_layers is None:
+            return super().propagate(graph)
+        for layer in self._model_layers:
+            graph = layer(graph)
+        return graph
+    
+    def get_config(self):
+        """Obtain model config."""
+        config = super().get_config()
+        if hasattr(self, '_model_layers') and self._model_layers is not None:
+            config['model_layers'] = [
+                keras.saving.serialize_keras_object(l) 
+                for l in self._model_layers
+            ]
+        return config 
+    
+    @classmethod
+    def from_config(cls, config: dict):
+        """Obtain model from model config."""
+        if 'model_layers' in config:
+            config['model_layers'] = [
+                keras.saving.deserialize_keras_object(l) 
+                for l in config['model_layers']
+            ]
+        return super().from_config(config)
+
+    def compile(
+        self, 
+        optimizer: keras.optimizers.Optimizer | str | None = 'rmsprop',
+        loss: keras.losses.Loss | str | None = None,
+        loss_weights: dict[str, float] = None,
+        metrics: list[keras.metrics.Metric] = None,
+        weighted_metrics: list[keras.metrics.Metric] | None = None,
+        run_eagerly: bool = False,
+        steps_per_execution: int = 1,
+        jit_compile: str | bool = False,
+        auto_scale_loss: bool = True,
+        **kwargs
+    ) -> None:
+        """Compiles the model.
+
+        Args:
+            optimizer:
+                The optimizer to be used (a `keras.optimizers.Optimizer` subclass).
+            loss:
+                The loss function to be used (a `keras.losses.Loss` subclass).
+            metrics:
+                A list of metrics to be used during training (`fit`) and evaluation
+                (`evaluate`). Should be `keras.metrics.Metric` subclasses.
+            kwargs:
+                See `Model.compile` in Keras documentation. 
+                May or may not apply here.
+        """
+        super().compile(
+            optimizer=optimizer,
+            loss=loss,
+            loss_weights=loss_weights,
+            metrics=metrics,
+            weighted_metrics=weighted_metrics,
+            run_eagerly=run_eagerly,
+            steps_per_execution=steps_per_execution,
+            jit_compile=jit_compile,
+            auto_scale_loss=auto_scale_loss,
+            **kwargs
+        )
+
+    def fit(self, x: tensors.GraphTensor | tf.data.Dataset, **kwargs):
+        """Fits the model.
+
+        Args:
+            x: 
+                A `GraphTensor` instance or a `tf.data.Dataset` constructed from
+                a `GraphTensor` instance. In comparison to a typical Keras model,
+                the label (typically denoted `y`) and the sample_weight (typically
+                denoted `sample_weight`) should be encoded in the context of the 
+                `GraphTensor` instance, as `label` and `weight` respectively.
+            validation_data:
+                A `GraphTensor` instance or a `tf.data.Dataset` constructed from
+                a `GraphTensor` instance. In comparison to a typical Keras model,
+                the label (typically denoted `y`) and the sample_weight (typically
+                denoted `sample_weight`) should be encoded in the context of the 
+                `GraphTensor` instance, as `label` and `weight` respectively.
+            validaton_split:
+                The fraction of training data to be used as validation data. 
+                Only works if a `GraphTensor` instance is passed as `x`. 
+            batch_size:
+                Number of samples per batch of computation.
+            epochs:
+                Number of iterations over the entire dataset.
+            callbacks:
+                A list of callbacks to apply during training.
+            kwargs:
+                See `Model.fit` in Keras documentation. 
+                May or may not apply here. 
+        """
+        batch_size = kwargs.get('batch_size', 32)
+        x_val = kwargs.pop('validation_data', None)
+        val_split = kwargs.pop('validation_split', None)
+        if x_val is not None and isinstance(x_val, tensors.GraphTensor):
+            x_val = _make_dataset(x_val, batch_size)
+        if isinstance(x, tensors.GraphTensor):
+            if val_split:
+                val_size = int(val_split * x.num_subgraphs)
+                x_val = _make_dataset(x[-val_size:], batch_size)
+                x = x[:-val_size]
+            x = _make_dataset(x, batch_size, shuffle=kwargs.get('shuffle', True))
+        return super().fit(x, validation_data=x_val, **kwargs)
+    
+    def evaluate(self, x: tensors.GraphTensor | tf.data.Dataset, **kwargs):
+        """Evaluation of the model.
+
+        Args:
+            x: 
+                A `GraphTensor` instance or a `tf.data.Dataset` constructed from
+                a `GraphTensor` instance. In comparison to a typical Keras model,
+                the label (typically denoted `y`) and the sample_weight (typically
+                denoted `sample_weight`) should be encoded in the context of the 
+                `GraphTensor` instance, as `label` and `weight` respectively.
+            batch_size:
+                Number of samples per batch of computation.
+            kwargs:
+                See `Model.evaluate` in Keras documentation. 
+                May or may not apply here. 
+        """
+        batch_size = kwargs.get('batch_size', 32)
+        if isinstance(x, tensors.GraphTensor):
+            x = _make_dataset(x, batch_size)
+        metric_results = super().evaluate(x, **kwargs)
+        return tf.nest.map_structure(lambda value: float(value), metric_results)
+    
+    def predict(self, x: tensors.GraphTensor | tf.data.Dataset, **kwargs):
+        """Makes predictions with the model.
+
+        Args:
+            x: 
+                A `GraphTensor` instance or a `tf.data.Dataset` constructed from
+                a `GraphTensor` instance.
+            batch_size:
+                Number of samples per batch of computation.
+            kwargs:
+                See `Model.predict` in Keras documentation. 
+                May or may not apply here. 
+        """
+        batch_size = kwargs.get('batch_size', 32)
+        if isinstance(x, tensors.GraphTensor):
+            x = _make_dataset(x, batch_size)
+        output = super().predict(x, **kwargs)
+        if tensors.is_graph(output):
+            return tensors.from_dict(output).flatten()
+        return output
+
+    def get_compile_config(self) -> dict | None:
+        config = super().get_compile_config()
+        if config is None:
+            return
+        return config
+
+    def compile_from_config(self, config: dict | None) -> None:
+        if config is None:
+            return
+        config = keras.utils.deserialize_keras_object(config)
+        self.compile(**config)
+        if hasattr(self, 'optimizer') and self.built:
+            self.optimizer.build(self.trainable_variables)
+
+    def save(
+        self,
+        filepath: str | Path,
+        *args,
+        **kwargs
+    ) -> None:
+        """Saves an entire model.
+        
+        Args:
+            filepath:
+                A string with the path to the model file (requires `.keras` suffix)
+        """
+        if not self.built:
+            raise ValueError('Cannot save model as it has not been built yet.')
+        super().save(filepath, *args, **kwargs)
+
+    @staticmethod
+    def load(
+        filepath: str | Path,
+        *args,
+        **kwargs
+    ) -> keras.Model:
+        """A `staticmethod` loading an entire model.
+
+        Args:
+            filepath:
+                A string with the path to the model file (requires `.keras` suffix)
+        """
+        return keras.models.load_model(filepath, *args, **kwargs)
+    
+    def save_weights(self, filepath, *args, **kwargs):
+        """Saves the weights of the model.
+
+        Args:
+            filepath:
+                A string with the path to the file (requires `.weights.h5` suffix)
+        """
+        path = Path(filepath).parent
+        path.mkdir(parents=True, exist_ok=True)
+        return super().save_weights(filepath, *args, **kwargs)
+
+    def load_weights(self, filepath, *args, **kwargs):
+        """Loads the weights from file saved via `save_weights()`.
+        
+        Args:
+            filepath:
+                A string with the path to the file (requires `.weights.h5` suffix)
+        """
+        super().load_weights(filepath, *args, **kwargs)
+
+    def embedding(self, layer_name: str = None) -> 'FunctionalGraphModel':
+        model = self
+        if not isinstance(model, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the embedding part of the model, '
+                'it needs to be a `FunctionalGraphModel`. '
+            )
+        inputs = model.input 
+        if not layer_name:
+            for layer in model.layers:
+                if isinstance(layer, layers.Readout):
+                    outputs = layer.output 
+        else:
+            layer = model.get_layer(layer_name)
+            outputs = (
+                layer.output if isinstance(layer, keras.layers.Layer) else None
+            )
+            if outputs is None:
+                raise ValueError(
+                    f'Could not find `{layer_name}` or '
+                    f'`{layer_name} is not a `keras.layers.Layer`.'
+                )
+        return self.__class__(inputs, outputs, name=f'{self.name}_embedding')
+
+    def backbone(self) -> 'FunctionalGraphModel':
+        if not isinstance(self, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the backbone part of the model, '
+                'it needs to be a `FunctionalGraphModel`, with a `Readout` '
+                'layer dividing the backbone and the head part of the model.'
+            )
+        inputs = self.input
+        outputs = None
+        for layer in self.layers:
+            if isinstance(layer, layers.Readout):
+                outputs = layer.output
+        if outputs is None:
+            raise ValueError(
+                'Could not extract output. `Readout` layer not found.'
+            )
+        return self.__class__(inputs, outputs, name=f'{self.name}_backbone')
+
+    def head(self) -> functional.Functional:
+        if not isinstance(self, FunctionalGraphModel):
+            raise ValueError(
+                'Currently, to extract the head part of the model, '
+                'it needs to be a `FunctionalGraphModel`, with a `Readout` '
+                'layer dividing the backbone and the head part of the model.'
+            )
+        inputs = None
+        for layer in self.layers:
+            if isinstance(layer, layers.Readout):
+                inputs = layer.output
+        if inputs is None:
+            raise ValueError(
+                'Could not extract input. `Readout` layer not found.'
+            )
+        outputs = layer.output
+        return keras.models.Model(inputs, outputs, name=f'{self.name}_head')
+
+    def train_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
+        with tf.GradientTape() as tape:
+            output = self(tensor, training=True)
+            y, y_pred, sample_weight = _get_loss_args(tensor, output)
+            loss = self.compute_loss(tensor, y, y_pred, sample_weight)
+            loss = self.optimizer.scale_loss(loss)
+        trainable_weights = self.trainable_weights 
+        gradients = tape.gradient(loss, trainable_weights)
+        self.optimizer.apply_gradients(zip(gradients, trainable_weights))
+        return self.compute_metrics(tensor, y, y_pred, sample_weight)
+    
+    def test_step(self, tensor: tensors.GraphTensor) -> dict[str, float]:
+        output = self(tensor, training=False)
+        y, y_pred, sample_weight = _get_loss_args(tensor, output)
+        return self.compute_metrics(tensor, y, y_pred, sample_weight)
+    
+    def predict_step(self, tensor: tensors.GraphTensor) -> np.ndarray:
+        output = self(tensor, training=False)
+        if tensors.is_graph(output):
+            if not isinstance(output, tensors.GraphTensor):
+                output = tensors.from_dict(output)
+            output = tensors.to_dict(output.unflatten())
+        return output
+
+    def compute_loss(self, x, y, y_pred, sample_weight=None):
+        return super().compute_loss(x, y, y_pred, sample_weight)
+        
+    def compute_metrics(self, x, y, y_pred, sample_weight=None) -> dict[str, float]:
+        loss = self.compute_loss(x, y, y_pred, sample_weight)
+        metric_results = {}
+        for metric in self.metrics:
+            if metric.name == "loss":
+                metric.update_state(loss)
+                metric_results[metric.name] = metric.result()
+            else:
+                metric.update_state(y, y_pred, sample_weight=sample_weight)
+                metric_results.update(metric.result())
+        return metric_results
+    
+
+@keras.saving.register_keras_serializable(package="molcraft")
+class FunctionalGraphModel(functional.Functional, GraphModel):
+
+    @property 
+    def layers(self):
+        return [
+            l for l in super().layers if not isinstance(l, keras.layers.InputLayer)
+        ]
+
+
+def save_model(model: GraphModel, filepath: str | Path, *args, **kwargs) -> None:
+    if not model.built:
+        raise ValueError(
+            'Model and its layers have not yet been (fully) built. '
+            'Build the model before saving it: `model.build(graph_spec)` '
+            'or `model(graph)`.'
+        )
+    keras.models.save_model(model, filepath, *args, **kwargs)
+
+def load_model(filepath: str | Path, inputs=None, *args, **kwargs) -> GraphModel:
+    return keras.models.load_model(filepath, *args, **kwargs)
+
+def create(
+    *layers: list[keras.layers.Layer],
+    **kwargs
+) -> GraphModel:
+    if isinstance(layers[0], list):
+        layers = layers[0]
+    return GraphModel.from_layers(
+        list(layers), **kwargs
+    )
+
+def interpret(
+    model: GraphModel,
+    graph_tensor: tensors.GraphTensor,
+) -> tensors.GraphTensor:
+    x = graph_tensor
+    if tensors.is_ragged(x):
+        x = x.flatten()
+    graph_indicator = x.graph_indicator
+    y_true = x.context.get('label')
+    features = []
+    with tf.GradientTape(watch_accessed_variables=False) as tape:
+        for layer in model.layers:
+            if isinstance(layer, keras.layers.InputLayer):
+                continue
+            if isinstance(layer, layers.GraphNetwork):
+                x, taped_features = layer.tape_propagate(x, tape, training=False)
+                features.extend(taped_features)
+            else:
+                if (
+                    isinstance(layer, layers.GraphConv) and 
+                    isinstance(x, tensors.GraphTensor)
+                ):
+                    tape.watch(x.node['feature'])
+                    features.append(x.node['feature'])
+                x = layer(x, training=False)
+        y_pred = x
+        if y_true is not None and len(y_true.shape) > 1: 
+            target = tf.gather_nd(y_pred, tf.where(y_true != 0))
+        else:
+            target = y_pred
+    gradients = tape.gradient(target, features)
+    features = keras.ops.concatenate(features, axis=-1)
+    gradients = keras.ops.concatenate(gradients, axis=-1)
+    alpha = ops.segment_mean(gradients, graph_indicator)
+    alpha = ops.gather(alpha, graph_indicator)
+    maps = keras.ops.where(gradients != 0, alpha * features, gradients)
+    maps = keras.ops.sum(maps, axis=-1)
+    return graph_tensor.update(
+        {
+            'node': {
+                'saliency': maps
+            }
+        }
+    )
+
+def saliency(
+    model: GraphModel,
+    graph_tensor: tensors.GraphTensor,
+) -> tensors.GraphTensor:
+    x = graph_tensor
+    if tensors.is_ragged(x):
+        x = x.flatten()
+    y_true = x.context.get('label')
+    with tf.GradientTape(watch_accessed_variables=False) as tape:
+        tape.watch(x.node['feature'])
+        y_pred = model(x, training=False)
+        if y_true is not None and len(y_true.shape) > 1: 
+            target = tf.gather_nd(y_pred, tf.where(y_true != 0))
+        else:
+            target = y_pred
+    gradients = tape.gradient(target, x.node['feature'])
+    gradients = keras.ops.absolute(gradients)
+    return graph_tensor.update(
+        {
+            'node': {
+                'feature_saliency': gradients
+            }
+        }
+    ) 
+    
+def _functional_init_arguments(args, kwargs):
+    return (
+        (len(args) == 2)
+        or (len(args) == 1 and "outputs" in kwargs)
+        or ("inputs" in kwargs and "outputs" in kwargs)
+    )
+
+def _make_dataset(x: tensors.GraphTensor, batch_size: int, shuffle: bool = False):
+    ds = tf.data.Dataset.from_tensor_slices(x)
+    if shuffle:
+        ds = ds.shuffle(buffer_size=ds.cardinality())
+    return ds.batch(batch_size).prefetch(-1)
+
+def _get_loss_args(
+    inputs: tensors.GraphTensor,
+    outputs: tensors.GraphTensor | tf.Tensor,
+) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor | None]:
+    if (
+        not isinstance(inputs, tensors.GraphTensor) and
+        tensors.is_graph(inputs)
+    ):
+        inputs = tensors.from_dict(inputs)
+    if (
+        not isinstance(outputs, tensors.GraphTensor) and
+        tensors.is_graph(outputs)
+    ):
+        outputs = tensors.from_dict(outputs)
+
+    if not isinstance(outputs, tensors.GraphTensor):
+        tensor, prediction = inputs, outputs
+    else:
+        tensor, prediction = outputs, None
+
+    if 'label' in tensor.context:
+        data = tensor.context
+    elif 'label' in tensor.node:
+        data = tensor.node
+    elif 'label' in tensor.edge:
+        data = tensor.edge
+    else:
+        raise ValueError(
+            'Could not find a `label` in the `GraphTensor`. Make sure a '
+            '`label` exists in either the `context`, `node` or `edge`.'
+        )
+
+    prediction = (
+        prediction if prediction is not None else data.get('prediction')
+    )
+    if prediction is None:
+        raise ValueError(
+            'Could not find a `prediction` in the `GraphTensor`. Make sure a '
+            '`prediction` exists in either the `context`, `node` or `edge`.'
+        )
+    return data['label'], prediction, data.get('sample_weight')