<commit_before>import warnings
import featureflow as ff
from preprocess import Preprocessor, PreprocessResult, Op


class PyTorchPreprocessResult(PreprocessResult):
    def __init__(self, data, op, inversion_data=None, inverse=None, name=None):
        super(PyTorchPreprocessResult, self).__init__(
            data, op, inversion_data, inverse, name)

    def __getstate__(self):
        forward_func = self.op._func
        inv_data_func = self.inversion_data._func
        backward_func = self.inverse._func
        network_params = self.op.network.state_dict()
        weights = dict(
            ((k, v.cpu().numpy()) for k, v in network_params.iteritems()))
        cls = self.op.network.__class__
        name = self.name

        kwargs = self.op.kwargs
        del kwargs['network']

        return dict(
            forward_func=forward_func,
            op_kwargs=kwargs,
            inv_data_func=inv_data_func,
            backward_func=backward_func,
            weights=weights,
            name=name,
            cls=cls)

    def __setstate__(self, state):
        import torch
        restored_weights = dict(
            ((k, torch.from_numpy(v).cuda())
             for k, v in state['weights'].iteritems()))

        network = state['cls']()
        network.load_state_dict(restored_weights)
        network.cuda()
        network.eval()
        self.op = Op(
            state['forward_func'], network=network, **state['op_kwargs'])
        self.inversion_data = Op(state['inv_data_func'],
                                 network=network)
        self.inverse = Op(state['backward_func'])
        self.name = state['name']

    def for_storage(self):
        return PyTorchPreprocessResult(
            None,
            self.op,
            self.inversion_data,
            self.inverse,
            self.name)


class PyTorchNetwork(Preprocessor):
    def __init__(self, trainer=None, needs=None):
        super(PyTorchNetwork, self).__init__(needs=needs)
        self.trainer = trainer
        self._cache = dict()

    def _forward_func(self):
        def x(d, network=None):
            from zounds.core import ArrayWithUnits, IdentityDimension
            from zounds.learn import apply_network

            result = apply_network(network, d, chunksize=500)
            try:
                return ArrayWithUnits(
                    result, d.dimensions[:-1] + (IdentityDimension(),))
            except AttributeError:
                return result
            except ValueError:
                # the number of dimensions has likely changed
                return result

        return x

    def _backward_func(self):
        def x(_):
            raise NotImplementedError()

        return x

    def _enqueue(self, data, pusher):
        k = self._dependency_name(pusher)
        self._cache[k] = data

    def _dequeue(self):
        if not self._finalized:
            raise ff.NotEnoughData()
        return self._cache

    def _process(self, data):
        data = self._extract_data(data)

        trained_network = self.trainer.train(data)

        try:
            forward_func = self._forward_func()
            processed_data = forward_func(data['data'], network=trained_network)
        except RuntimeError as e:
            processed_data = None
            # the dataset may be too large to fit onto the GPU all at once
            warnings.warn(e.message)
        except KeyError:
            processed_data = None

        op = self.transform(network=trained_network)
        inv_data = self.inversion_data()
        inv = self.inverse_transform()

        yield PyTorchPreprocessResult(
            processed_data,
            op,
            inversion_data=inv_data,
            inverse=inv,
            name='PyTorchNetwork')


class PyTorchGan(PyTorchNetwork):
    def __init__(self, apply_network='generator', trainer=None, needs=None):
        super(PyTorchGan, self).__init__(trainer=trainer, needs=needs)

        if apply_network not in ('generator', 'discriminator'):
            raise ValueError(
                'apply_network must be one of (generator, discriminator)')

        self.apply_network = apply_network
        self._cache = None

    def _forward_func(self):
        def x(d, network=None, apply_network=None):
            import torch
            from torch.autograd import Variable
            import numpy as np
            from zounds.core import ArrayWithUnits, IdentityDimension

            if apply_network == 'generator':
                n = network.generator
            else:
                n = network.discriminator

            chunks = []
            batch_size = 128

            for i in xrange(0, len(d), batch_size):
                tensor = torch.from_numpy(
                    d[i:i + batch_size].astype(np.float32))
                gpu = tensor.cuda()
                v = Variable(gpu)
                chunks.append(n(v).data.cpu().numpy())

            result = np.concatenate(chunks)

            # tensor = torch.from_numpy(d.astype(np.float32))
            # gpu = tensor.cuda()
            # v = Variable(gpu)
            # result = n(v).data.cpu().numpy()

            try:<commit_msg>Remove some unused, commented code<commit_after>import warnings
import featureflow as ff
from preprocess import Preprocessor, PreprocessResult, Op


class PyTorchPreprocessResult(PreprocessResult):
    def __init__(self, data, op, inversion_data=None, inverse=None, name=None):
        super(PyTorchPreprocessResult, self).__init__(
            data, op, inversion_data, inverse, name)

    def __getstate__(self):
        forward_func = self.op._func
        inv_data_func = self.inversion_data._func
        backward_func = self.inverse._func
        network_params = self.op.network.state_dict()
        weights = dict(
            ((k, v.cpu().numpy()) for k, v in network_params.iteritems()))
        cls = self.op.network.__class__
        name = self.name

        kwargs = self.op.kwargs
        del kwargs['network']

        return dict(
            forward_func=forward_func,
            op_kwargs=kwargs,
            inv_data_func=inv_data_func,
            backward_func=backward_func,
            weights=weights,
            name=name,
            cls=cls)

    def __setstate__(self, state):
        import torch
        restored_weights = dict(
            ((k, torch.from_numpy(v).cuda())
             for k, v in state['weights'].iteritems()))

        network = state['cls']()
        network.load_state_dict(restored_weights)
        network.cuda()
        network.eval()
        self.op = Op(
            state['forward_func'], network=network, **state['op_kwargs'])
        self.inversion_data = Op(state['inv_data_func'],
                                 network=network)
        self.inverse = Op(state['backward_func'])
        self.name = state['name']

    def for_storage(self):
        return PyTorchPreprocessResult(
            None,
            self.op,
            self.inversion_data,
            self.inverse,
            self.name)


class PyTorchNetwork(Preprocessor):
    def __init__(self, trainer=None, needs=None):
        super(PyTorchNetwork, self).__init__(needs=needs)
        self.trainer = trainer
        self._cache = dict()

    def _forward_func(self):
        def x(d, network=None):
            from zounds.core import ArrayWithUnits, IdentityDimension
            from zounds.learn import apply_network

            result = apply_network(network, d, chunksize=500)
            try:
                return ArrayWithUnits(
                    result, d.dimensions[:-1] + (IdentityDimension(),))
            except AttributeError:
                return result
            except ValueError:
                # the number of dimensions has likely changed
                return result

        return x

    def _backward_func(self):
        def x(_):
            raise NotImplementedError()

        return x

    def _enqueue(self, data, pusher):
        k = self._dependency_name(pusher)
        self._cache[k] = data

    def _dequeue(self):
        if not self._finalized:
            raise ff.NotEnoughData()
        return self._cache

    def _process(self, data):
        data = self._extract_data(data)

        trained_network = self.trainer.train(data)

        try:
            forward_func = self._forward_func()
            processed_data = forward_func(data['data'], network=trained_network)
        except RuntimeError as e:
            processed_data = None
            # the dataset may be too large to fit onto the GPU all at once
            warnings.warn(e.message)
        except KeyError:
            processed_data = None

        op = self.transform(network=trained_network)
        inv_data = self.inversion_data()
        inv = self.inverse_transform()

        yield PyTorchPreprocessResult(
            processed_data,
            op,
            inversion_data=inv_data,
            inverse=inv,
            name='PyTorchNetwork')


class PyTorchGan(PyTorchNetwork):
    def __init__(self, apply_network='generator', trainer=None, needs=None):
        super(PyTorchGan, self).__init__(trainer=trainer, needs=needs)

        if apply_network not in ('generator', 'discriminator'):
            raise ValueError(
                'apply_network must be one of (generator, discriminator)')

        self.apply_network = apply_network
        self._cache = None

    def _forward_func(self):
        def x(d, network=None, apply_network=None):
            import torch
            from torch.autograd import Variable
            import numpy as np
            from zounds.core import ArrayWithUnits, IdentityDimension

            if apply_network == 'generator':
                n = network.generator
            else:
                n = network.discriminator

            chunks = []
            batch_size = 128

            for i in xrange(0, len(d), batch_size):
                tensor = torch.from_numpy(
                    d[i:i + batch_size].astype(np.float32))
                gpu = tensor.cuda()
                v = Variable(gpu)
                chunks.append(n(v).data.cpu().numpy())

            result = np.concatenate(chunks)

            try:
{'fraction_non_alphanumeric': 0.09074019751637821, 'fraction_numerical': 0.002151168475603794, 'mean_word_length': 1.9298195359495847, '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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3196303', 'n_tokens_mistral': 2851, 'n_tokens_neox': 2755, 'n_words': 635}