<commit_before>from .base import Model
from .base import DifferentiableModel


class ModelWrapper(Model):
    """Base class for models that wrap other models.

    This base class can be used to implement model wrappers
    that turn models into new models, for example by preprocessing
    the input or modifying the gradient.

    Parameters
    ----------
    model : :class:`Model`
        The model that is wrapped.

    """

    def __init__(self, model):
        super(ModelWrapper, self).__init__(
            bounds=model.bounds(),
            channel_axis=model.channel_axis())

        self.wrapped_model = model

    def __enter__(self):
        assert self.wrapped_model.__enter__() == self.wrapped_model
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return self.wrapped_model.__exit__(exc_type, exc_value, traceback)

    def batch_predictions(self, images):
        return self.wrapped_model.batch_predictions(images)

    def predictions(self, image):
        return self.wrapped_model.predictions(image)

    def num_classes(self):
        return self.wrapped_model.num_classes()


class DifferentiableModelWrapper(ModelWrapper):
    """Base class for models that wrap other models and provide
    gradient methods.

    This base class can be used to implement model wrappers
    that turn models into new models, for example by preprocessing
    the input or modifying the gradient.

    Parameters
    ----------
    model : :class:`Model`
        The model that is wrapped.

    """

    def predictions_and_gradient(self, image, label):
        return self.wrapped_model.predictions_and_gradient(image)

    def gradient(self, image, label):
        return self.wrapped_model.gradient(image)

    def backward(self, gradient, image):
        return self.wrapped_model.backward(gradient, image)

<commit_msg>bug fixes<commit_after>from .base import Model
from .base import DifferentiableModel


class ModelWrapper(Model):
    """Base class for models that wrap other models.

    This base class can be used to implement model wrappers
    that turn models into new models, for example by preprocessing
    the input or modifying the gradient.

    Parameters
    ----------
    model : :class:`Model`
        The model that is wrapped.

    """

    def __init__(self, model):
        super(ModelWrapper, self).__init__(
            bounds=model.bounds(),
            channel_axis=model.channel_axis())

        self.wrapped_model = model

    def __enter__(self):
        assert self.wrapped_model.__enter__() == self.wrapped_model
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return self.wrapped_model.__exit__(exc_type, exc_value, traceback)

    def batch_predictions(self, images):
        return self.wrapped_model.batch_predictions(images)

    def predictions(self, image):
        return self.wrapped_model.predictions(image)

    def num_classes(self):
        return self.wrapped_model.num_classes()


class DifferentiableModelWrapper(ModelWrapper):
    """Base class for models that wrap other models and provide
    gradient methods.

    This base class can be used to implement model wrappers
    that turn models into new models, for example by preprocessing
    the input or modifying the gradient.

    Parameters
    ----------
    model : :class:`Model`
        The model that is wrapped.

    """

    def predictions_and_gradient(self, image, label):
        return self.wrapped_model.predictions_and_gradient(image, label)

    def gradient(self, image, label):
        return self.wrapped_model.gradient(image, label)

    def backward(self, gradient, image):
        return self.wrapped_model.backward(gradient, image)


{'fraction_non_alphanumeric': 0.08328884143085959, 'fraction_numerical': 0.0, 'mean_word_length': 2.9861702127659573, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2820307', 'n_tokens_mistral': 1025, 'n_tokens_neox': 974, 'n_words': 319}