<commit_before># pylint: disable=undefined-variable
""" Contains base class for tensorflow models """

import os
import glob
import re
import json
import numpy as np
import tensorflow as tf

from ..base import BaseModel
from .layers import mip, conv_block
from .losses import dice
from ...utils import copy1


LOSSES = {
    'mse': tf.losses.mean_squared_error,
    'ce': tf.losses.softmax_cross_entropy,
    'crossentropy': tf.losses.softmax_cross_entropy,
    'absolutedifference': tf.losses.absolute_difference,
    'l1': tf.losses.absolute_difference,
    'cosine': tf.losses.cosine_distance,
    'cos': tf.losses.cosine_distance,
    'hinge': tf.losses.hinge_loss,
    'huber': tf.losses.huber_loss,
    'logloss': tf.losses.log_loss,
    'dice': dice
}

DECAYS = {
    'exp': tf.train.exponential_decay,
    'invtime': tf.train.inverse_time_decay,
    'naturalexp': tf.train.natural_exp_decay,
    'const': tf.train.piecewise_constant,
    'poly': tf.train.polynomial_decay
}



class TFModel(BaseModel):
    r""" Base class for all tensorflow models

    Attributes
    ----------
    session : tf.Session

    graph : tf.Graph

    is_training : tf.Tensor

    global_step : tf.Tensor

    loss : tf.Tensor

    train_step : tf.Operation

    from BaseModel:

    name : str - a model name

    config : dict - configuration parameters


    **Configuration**

    ``build`` and ``load`` are inherited from :class:`.BaseModel`.

    session : dict - `Tensorflow session parameters <https://www.tensorflow.org/api_docs/python/tf/Session#__init__>`_.

    inputs : dict - model inputs (see :meth:`._make_inputs`)

    loss - a loss function, might be one of:
        - short name (`'mse'`, `'ce'`, `'l1'`, `'cos'`, `'hinge'`, `'huber'`, `'logloss'`, `'dice'`)
        - a function name from `tf.losses <https://www.tensorflow.org/api_docs/python/tf/losses>`_
          (e.g. `'absolute_difference'` or `'sparse_softmax_cross_entropy'`)
        - callable

        Examples:

        - ``{'loss': 'mse'}``
        - ``{'loss': 'sigmoid_cross_entropy'}``
        - ``{'loss': tf.losses.huber_loss}``
        - ``{'loss': external_loss_fn}``

    decay - a learning rate decay algorithm might be defined in one of three formats:
        - name
        - tuple (name, args)
        - dict {'name': name, \**args}

        where name might be one of:

        - short name ('exp', 'invtime', 'naturalexp', 'const', 'poly')
        - a function name from `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_
          (e.g. 'exponential_decay')
        - a callable

        Examples:

        - ``{'decay': 'exp'}``
        - ``{'decay': ('polynomial_decay', {'decay_steps':10000})}``
        - ``{'decay': {'name': tf.train.inverse_time_decay, 'decay_rate': .5}``

    optimizer - an optimizer might be defined in one of three formats:
            - name
            - tuple (name, args)
            - dict {'name': name, \**args}

            where name might be one of:

            - short name (e.g. 'Adam', 'Adagrad', any optimiser from
              `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_ without a word `Optimiser`)
            - a function name from `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_
              (e.g. 'FtlrOptimizer')
            - a callable

        Examples:

        - ``{'optimizer': 'Adam'}``
        - ``{'optimizer': ('Ftlr', {'learning_rate_power': 0})}``
        - ``{'optimizer': {'name': 'Adagrad', 'initial_accumulator_value': 0.01}``
        - ``{'optimizer': functools.partial(tf.train.MomentumOptimizer, momentum=0.95)}``
        - ``{'optimizer': some_optimizer_fn}``
    """

    def __init__(self, *args, **kwargs):
        self.session = kwargs.get('session', None)
        self.graph = tf.Graph() if self.session is None else self.session.graph
        self._graph_context = None
        self.is_training = None
        self.global_step = None
        self.loss = None
        self.train_step = None
        self._attrs = []
        self._to_classes = {}
        self._inputs = {}

        super().__init__(*args, **kwargs)


    def __enter__(self):
        """ Enter the model graph context """
        self._graph_context = self.graph.as_default()
        return self._graph_context.__enter__()

    def __exit__(self, exception_type, exception_value, exception_traceback):
        """ Exit the model graph context """
        return self._graph_context.__exit__(exception_type, exception_value, exception_traceback)

    def _build(self):
        """ Define a model architecture

        It takes just 2 steps:

        #. Define names for all placeholders and make input tensors by calling :meth:`._make_inputs`::

            def _build():
                names = ['images', 'labels']
                placeholders, inputs = self._make_inputs(names)

           If config does not contain any name from ``names``, :exc:`KeyError` is raised.

        #. Add your layers.

        And that's pretty it.

        Things worth mentioning:

        #. Input data and its parameters should be defined in configuration under ``inputs`` key.

        #. :meth:`._make_inputs` returns ``placeholders`` and ``inputs`` which are dicts
           with the same keys as ``config['inputs']``.
           They contain all placeholders and tensors after reshaping/transformations, correspondingly.
           Use ``inputs`` to build a model. While placeholder names will also be needed later in
           ``train`` and ``predict`` methods.

        #. You might want to use a convenient multidimensional :func:`~.layers.conv_block`,
           as well as :func:`~.layers.global_average_pooling`,
           :func:`~.layers.mip`, or other predefined layers.
           Of course, you can use usual `tensorflow layers <https://www.tensorflow.org/api_docs/python/tf/layers>`_.

        #. For dropout, batch norm, etc you might use a predefined ``self.is_training`` tensor.

        #. For decay and training control you might use a predefined ``self.global_step`` tensor.

        #. In many cases there is no need to write a loss function, learning rate decay and optimizer
           as they might be defined through config.

        #. For a configured loss one of the inputs should have a name ``targets`` and
           one of the tensors in your model should have a name ``predictions``.
           They will be used in a loss function.

        #. If you have defined your own loss function, call `tf.losses.add_loss(...)
           <https://www.tensorflow.org/api_docs/python/tf/losses/add_loss>`_.

        #. If you need to use your own optimizer, assign ``self.train_step`` to the train step operation.
           Don't forget about `UPDATE_OPS control dependency
           <https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization>`_.

        Notes
        -----
        This method is executed within a self.graph context
        """
        input_names = []
        placeholders, inputs = self._make_inputs(input_names)
        _ = placeholders, inputs


    def build(self, *args, **kwargs):
        """ Build the model

        #. Create a graph
        #. Define ``is_training`` and ``global_step`` tensors
        #. Define a model architecture by calling :meth:``._build``
        #. Create a loss function from config
        #. Create an optimizer and define a train step from config
        #. `Set UPDATE_OPS control dependency on train step
           <https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization>`_
        #. Create a tensorflow session
        """
        with self.graph.as_default():
            self.store_to_attr('is_training', tf.placeholder(tf.bool, name='is_training'))
            self.store_to_attr('global_step', tf.Variable(0, trainable=False, name='global_step'))

            self._build()

            self._make_loss()
            self.store_to_attr('loss', tf.losses.get_total_loss())

            if self.train_step is None:
                optimizer = self._make_optimizer()

                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                with tf.control_dependencies(update_ops):
                    self.store_to_attr('train_step', optimizer.minimize(self.loss, global_step=self.global_step))
            else:
                self.store_to_attr('train_step', self.train_step)

            session_config = self.get_from_config('session', {})
            self.session = tf.Session(**session_config)
            self.session.run(tf.global_variables_initializer())

    def _make_inputs(self, names=None):
        """ Make model input data using config

        In the config's inputs section it looks for ``names``, creates placeholders required, and
        makes some typical transformations (like one-hot-encoding and reshaping), if needed.

        **Configuration**

        inputs : dict
            - key : str
                a placeholder name
            - values : dict or tuple
                each input's config

        Input config:

        ``dtype`` : str or tf.DType (by default 'float32')
            data type

        ``shape`` : int, tuple, list or None (default)
            a desired tensor shape which includes the number of channels/classes and doesn't include a batch size.

        ``classes`` : array-like or None (default)
            an array of class labels if data labels are strings or anything else except ``np.arange(num_classes)``

        ``data_format`` : str {``'channels_first'``, ``'channels_last'``} or {``'f'``, ``'l'``}
            The ordering of the dimensions in the inputs. Default is 'channels_last'.

        ``transform`` : str or callable
            Predefined transforms are

            - ``'ohe'`` - one-hot encoding
            - ``'mip @ d'`` - maximum intensity projection :func:`~.layers.mip`
              with depth ``d`` (should be int)

        ``name`` : str
            a name for the transformed and reshaped tensor.

        If an input config is a tuple, it should contain all items exactly in the order shown above:
        dtype, shape, classes, data_format, transform, name.
        If an item is None, the default value will be used instead.

        **How it works**

        Input data will be cast to ``dtype``, then reshaped to ``shape`` in accordance with ``data_format``
        and finally transformed with a ``transform`` function. The resulting tensor will have the name ``name``.
        The original placeholder tensor will have the name ``key``.

        Parameters
        ----------
        names : list
            placeholder names that are expected in the config inputs section

        Raises
        ------
        KeyError if there is any name missing in the config's input section.
        ValueError if there are duplicate names.

        Returns
        -------
        placeholders : dict
            key : str
                a placeholder name
            value : tf.Tensor
                placeholder tensor
        tensors : dict
            key : str
                a placeholder name
            value : tf.Tensor
                placeholder tensor after reshaping and transformations
        """
        config = self.get_from_config('inputs') or {}
        config = copy1(config)

        names = names or []
        missing_names = set(names) - set(config.keys())
        if len(missing_names) > 0:
            raise KeyError("Inputs should contain {} names".format(missing_names))

        placeholder_names = set(config.keys())
        tensor_names = set(x.get('name') for x in config.values() if x.get('name'))
        wrong_names = placeholder_names & tensor_names
        if len(wrong_names) > 0:
            raise ValueError('Inputs contain duplicate names:', wrong_names)

        param_names = ('dtype', 'shape', 'classes', 'data_format', 'transform', 'name')
        defaults = dict(data_format='channels_last')

        placeholders = dict()
        tensors = dict()

        for input_name, input_config in config.items():
            if isinstance(input_config, (tuple, list)):
                input_config = list(input_config) + [None for _ in param_names]
                input_config = input_config[:len(param_names)]
                input_config = dict(zip(param_names, input_config))
                input_config = dict((k, v) for k, v in input_config.items() if v is not None)
            input_config = {**defaults, **input_config}

            shape = input_config.get('shape')
            if isinstance(shape, int):
                input_config['shape'] = (shape,)

            self._inputs[input_name] = dict(config=input_config)

            if self.has_classes(input_name):
                dtype = input_config.get('dtype', tf.int32)
            else:
                dtype = input_config.get('dtype', 'float')
            tensor = tf.placeholder(dtype, name=input_name)
            placeholders[input_name] = tensor

            if input_config.get('data_format') == 'l':
                input_config['data_format'] = 'channels_last'
            elif input_config.get('data_format') == 'f':
                input_config['data_format'] = 'channels_first'

            self._inputs[input_name] = dict(config=input_config)
            tensor = self._make_transform(input_name, tensor, input_config)

            shape = input_config.get('shape')
            if input_config.get('__reshape', True) and isinstance(shape, (list, tuple)):
                input_config['shape'] = shape
                tensor = tf.reshape(tensor, [-1] + list(shape))

            name = input_config.get('name')
            if name is not None:
                tensor = tf.identity(tensor, name=name)

            tensors[input_name] = tensor

            self._inputs[input_name] = dict(config=input_config, placeholder=placeholders[input_name], tensor=tensor)

        return placeholders, tensors

    def _make_transform(self, input_name, tensor, config):
        if config is not None:
            transform_name = config.get('transform')
            if isinstance(transform_name, str):
                transforms = {
                    'ohe': self._make_ohe,
                    'mip': self._make_mip
                }

                if transform_name.startswith('mip'):
                    parts = transform_name.split('@')
                    transform_name = parts[0].strip()
                    config['depth'] = int(parts[1])
                    # mip has to know shape, so we need call reshape first
                    shape = config.get('shape')
                    tensor = tf.reshape(tensor, [-1] + list(shape))
                    # do not make reshape twice
                    config.update({'__reshape': False})

                tensor = transforms[transform_name](input_name, tensor, config)
            elif callable(transform_name):
                tensor = transform_name(tensor)
            elif transform_name is not None:
                raise ValueError("Unknown transform {}".format(transform_name))
        return tensor

    def _make_ohe(self, input_name, tensor, config):
        if config.get('shape') is None and config.get('classes') is None:
            raise ValueError("shape and classes cannot be both None for input '{}'".format(input_name))

        num_classes = self.num_classes(input_name)
        axis = -1 if self.data_format(input_name) else 1
        tensor = tf.one_hot(tensor, depth=num_classes, axis=axis)
        return tensor

    def to_classes(self, tensor, input_name, name=None):
        """ Convert tensor with labels to classes of ``input_name`` """
        if tensor.dtype in [tf.float16, tf.float32, tf.float64]:
            tensor = tf.argmax(tensor, axis=-1, name=name)
        if self.has_classes(input_name):
            self._to_classes.update({tensor: input_name})
        return tensor

    def _make_mip(self, input_name, tensor, config):
        depth = config.get('depth')
        if depth is None:
            raise ValueError("mip should be specifies as mip @ depth")
        tensor = mip(tensor, depth=depth, data_format=self.data_format(input_name))
        return tensor

    def _unpack_fn_from_config(self, param, default=None):
        par = self.get_from_config(param, default)

        if par is None:
            return None, None

        if isinstance(par, (tuple, list)):
            if len(par) == 0:
                par_name = None
            elif len(par) == 1:
                par_name, par_args = par[0], {}
            elif len(par) == 2:
                par_name, par_args = par
            else:
                par_name, par_args = par[0], par[1:]
        elif isinstance(par, dict):
            par = par.copy()
            par_name, par_args = par.pop('name', None), par
        else:
            par_name, par_args = par, {}

        return par_name, par_args

    def _make_loss(self):
        """ Return a loss function from config """
        loss = self.get_from_config("loss")
        if loss is None:
            if len(tf.losses.get_losses()) == 0:
                raise ValueError("Loss is not defined in the model %s" % self)
        elif isinstance(loss, str) and hasattr(tf.losses, loss):
            loss = getattr(tf.losses, loss)
        elif isinstance(loss, str):
            loss = LOSSES.get(re.sub('[-_ ]', '', loss).lower(), None)
        elif callable(loss):
            pass
        else:
            raise ValueError("Unknown loss", loss)

        if loss is not None:
            try:
                predictions = self.graph.get_tensor_by_name("predictions:0")
                targets = self.graph.get_tensor_by_name("targets:0")
            except KeyError:
                raise KeyError("Model %s does not have 'predictions' or 'targets' tensors" % self.name)
            else:
                tf.losses.add_loss(loss(targets, predictions))

    def _make_decay(self):
        decay_name, decay_args = self._unpack_fn_from_config('decay')

        if decay_name is None:
            pass
        elif callable(decay_name):
            pass
        elif isinstance(decay_name, str) and hasattr(tf.train, decay_name):
            decay_name = getattr(tf.train, decay_name)
        elif decay_name in DECAYS:
            decay_name = DECAYS.get(re.sub('[-_ ]', '', decay_name).lower(), None)
        else:
            raise ValueError("Unknown learning rate decay method", decay_name)

        return decay_name, decay_args

    def _make_optimizer(self):
        optimizer_name, optimizer_args = self._unpack_fn_from_config('optimizer', 'Adam')

        if callable(optimizer_name):
            pass
        elif isinstance(optimizer_name, str) and hasattr(tf.train, optimizer_name):
            optimizer_name = getattr(tf.train, optimizer_name)
        elif isinstance(optimizer_name, str) and hasattr(tf.train, optimizer_name + 'Optimizer'):
            optimizer_name = getattr(tf.train, optimizer_name + 'Optimizer')
        else:
            raise ValueError("Unknown optimizer", optimizer_name)

        decay_name, decay_args = self._make_decay()
        if decay_name is not None:
            optimizer_args['learning_rate'] = decay_name(**decay_args, global_step=self.global_step)

        optimizer = optimizer_name(**optimizer_args)

        return optimizer

    def get_number_of_trainable_vars(self):
        """ Return the number of trainable variable in the model graph """
        with self.graph:
            arr = np.asarray([np.prod(self.get_shape(v)) for v in tf.trainable_variables()])
        return np.sum(arr)

    def get_tensor_config(self, tensor, **kwargs):
        """ Return tensor configuration

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        dict
            tensor config (see :meth:`._make_inputs`)

        Raises
        ------
        ValueError shape in tensor configuration isn't int, tuple or list
        """
        if isinstance(tensor, tf.Tensor):
            names = [n for n, i in self._inputs.items() if tensor in [i['placeholder'], i['tensor']]]
            if len(names) > 0:
                input_name = names[0]
            else:
                input_name = tensor.name
        elif isinstance(tensor, str):
            if tensor in self._inputs:
                input_name = tensor
            else:
                input_name = self._map_name(tensor)
        else:
            raise TypeError("Tensor can be tf.Tensor or string, but given %s" % type(tensor))

        if input_name in self._inputs:
            config = self._inputs[input_name]['config']
            shape = config.get('shape')
            if isinstance(shape, int):
                shape = (shape,)
            kwargs['shape'] = shape

            if 'classes' not in config and 'shape' not in config:
                raise ValueError("one of 'classes' or 'shape' must be present in inputs config")
            if 'classes' not in config:
                if not isinstance(shape, (list, tuple)):
                    raise ValueError('shape must be int, tuple or list but {} was given'.format(type(shape)))
        else:
            tensor = self.graph.get_tensor_by_name(input_name)
            shape = tensor.get_shape().as_list()[1:]
            config = dict(dtype=tensor.dtype, shape=shape, name=tensor.name, data_format='channels_last')
            config = {**config, **kwargs}

        return config


    def num_channels(self, tensor, **kwargs):
        """ Return the number of channels in the tensor

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        number of channels : int
        """
        config = self.get_tensor_config(tensor, **kwargs)
        shape = config.get('shape')
        data_format = config.get('data_format')
        channels_dim = -1 if data_format == "channels_last" or not data_format.startswith("NC") else 0
        return shape[channels_dim]

    def has_classes(self, tensor):
        """ Check if a tensor has classes defined in the config """
        config = self.get_tensor_config(tensor)
        has = config.get('classes') is not None
        return has

    def classes(self, tensor):
        """ Return the  number of classes """
        config = self.get_tensor_config(tensor)
        return config.get('classes')

    def num_classes(self, tensor):
        """ Return the  number of classes """
        if self.has_classes(tensor):
            return len(self.classes(tensor))
        return self.num_channels(tensor)

    def spatial_dim(self, tensor, **kwargs):
        """ Return the tensor spatial dimensionality (without channels dimension)

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        number of spatial dimensions : int

        Raises
        ------
        ValueError shape in tensor configuration isn't int, tuple or list
        """
        config = self.get_tensor_config(tensor, **kwargs)
        return config.get('shape') - 1

    def data_format(self, tensor, **kwargs):
        """ Return the tensor data format (channels_last or channels_first)

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        data_format : str
        """
        config = self.get_tensor_config(tensor, **kwargs)
        return config.get('data_format')<commit_msg>fix spatial_dim<commit_after># pylint: disable=undefined-variable
""" Contains base class for tensorflow models """

import os
import glob
import re
import json
import numpy as np
import tensorflow as tf

from ..base import BaseModel
from .layers import mip, conv_block
from .losses import dice
from ...utils import copy1


LOSSES = {
    'mse': tf.losses.mean_squared_error,
    'ce': tf.losses.softmax_cross_entropy,
    'crossentropy': tf.losses.softmax_cross_entropy,
    'absolutedifference': tf.losses.absolute_difference,
    'l1': tf.losses.absolute_difference,
    'cosine': tf.losses.cosine_distance,
    'cos': tf.losses.cosine_distance,
    'hinge': tf.losses.hinge_loss,
    'huber': tf.losses.huber_loss,
    'logloss': tf.losses.log_loss,
    'dice': dice
}

DECAYS = {
    'exp': tf.train.exponential_decay,
    'invtime': tf.train.inverse_time_decay,
    'naturalexp': tf.train.natural_exp_decay,
    'const': tf.train.piecewise_constant,
    'poly': tf.train.polynomial_decay
}



class TFModel(BaseModel):
    r""" Base class for all tensorflow models

    Attributes
    ----------
    session : tf.Session

    graph : tf.Graph

    is_training : tf.Tensor

    global_step : tf.Tensor

    loss : tf.Tensor

    train_step : tf.Operation

    from BaseModel:

    name : str - a model name

    config : dict - configuration parameters


    **Configuration**

    ``build`` and ``load`` are inherited from :class:`.BaseModel`.

    session : dict - `Tensorflow session parameters <https://www.tensorflow.org/api_docs/python/tf/Session#__init__>`_.

    inputs : dict - model inputs (see :meth:`._make_inputs`)

    loss - a loss function, might be one of:
        - short name (`'mse'`, `'ce'`, `'l1'`, `'cos'`, `'hinge'`, `'huber'`, `'logloss'`, `'dice'`)
        - a function name from `tf.losses <https://www.tensorflow.org/api_docs/python/tf/losses>`_
          (e.g. `'absolute_difference'` or `'sparse_softmax_cross_entropy'`)
        - callable

        Examples:

        - ``{'loss': 'mse'}``
        - ``{'loss': 'sigmoid_cross_entropy'}``
        - ``{'loss': tf.losses.huber_loss}``
        - ``{'loss': external_loss_fn}``

    decay - a learning rate decay algorithm might be defined in one of three formats:
        - name
        - tuple (name, args)
        - dict {'name': name, \**args}

        where name might be one of:

        - short name ('exp', 'invtime', 'naturalexp', 'const', 'poly')
        - a function name from `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_
          (e.g. 'exponential_decay')
        - a callable

        Examples:

        - ``{'decay': 'exp'}``
        - ``{'decay': ('polynomial_decay', {'decay_steps':10000})}``
        - ``{'decay': {'name': tf.train.inverse_time_decay, 'decay_rate': .5}``

    optimizer - an optimizer might be defined in one of three formats:
            - name
            - tuple (name, args)
            - dict {'name': name, \**args}

            where name might be one of:

            - short name (e.g. 'Adam', 'Adagrad', any optimiser from
              `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_ without a word `Optimiser`)
            - a function name from `tf.train <https://www.tensorflow.org/api_docs/python/tf/train>`_
              (e.g. 'FtlrOptimizer')
            - a callable

        Examples:

        - ``{'optimizer': 'Adam'}``
        - ``{'optimizer': ('Ftlr', {'learning_rate_power': 0})}``
        - ``{'optimizer': {'name': 'Adagrad', 'initial_accumulator_value': 0.01}``
        - ``{'optimizer': functools.partial(tf.train.MomentumOptimizer, momentum=0.95)}``
        - ``{'optimizer': some_optimizer_fn}``
    """

    def __init__(self, *args, **kwargs):
        self.session = kwargs.get('session', None)
        self.graph = tf.Graph() if self.session is None else self.session.graph
        self._graph_context = None
        self.is_training = None
        self.global_step = None
        self.loss = None
        self.train_step = None
        self._attrs = []
        self._to_classes = {}
        self._inputs = {}

        super().__init__(*args, **kwargs)


    def __enter__(self):
        """ Enter the model graph context """
        self._graph_context = self.graph.as_default()
        return self._graph_context.__enter__()

    def __exit__(self, exception_type, exception_value, exception_traceback):
        """ Exit the model graph context """
        return self._graph_context.__exit__(exception_type, exception_value, exception_traceback)

    def _build(self):
        """ Define a model architecture

        It takes just 2 steps:

        #. Define names for all placeholders and make input tensors by calling :meth:`._make_inputs`::

            def _build():
                names = ['images', 'labels']
                placeholders, inputs = self._make_inputs(names)

           If config does not contain any name from ``names``, :exc:`KeyError` is raised.

        #. Add your layers.

        And that's pretty it.

        Things worth mentioning:

        #. Input data and its parameters should be defined in configuration under ``inputs`` key.

        #. :meth:`._make_inputs` returns ``placeholders`` and ``inputs`` which are dicts
           with the same keys as ``config['inputs']``.
           They contain all placeholders and tensors after reshaping/transformations, correspondingly.
           Use ``inputs`` to build a model. While placeholder names will also be needed later in
           ``train`` and ``predict`` methods.

        #. You might want to use a convenient multidimensional :func:`~.layers.conv_block`,
           as well as :func:`~.layers.global_average_pooling`,
           :func:`~.layers.mip`, or other predefined layers.
           Of course, you can use usual `tensorflow layers <https://www.tensorflow.org/api_docs/python/tf/layers>`_.

        #. For dropout, batch norm, etc you might use a predefined ``self.is_training`` tensor.

        #. For decay and training control you might use a predefined ``self.global_step`` tensor.

        #. In many cases there is no need to write a loss function, learning rate decay and optimizer
           as they might be defined through config.

        #. For a configured loss one of the inputs should have a name ``targets`` and
           one of the tensors in your model should have a name ``predictions``.
           They will be used in a loss function.

        #. If you have defined your own loss function, call `tf.losses.add_loss(...)
           <https://www.tensorflow.org/api_docs/python/tf/losses/add_loss>`_.

        #. If you need to use your own optimizer, assign ``self.train_step`` to the train step operation.
           Don't forget about `UPDATE_OPS control dependency
           <https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization>`_.

        Notes
        -----
        This method is executed within a self.graph context
        """
        input_names = []
        placeholders, inputs = self._make_inputs(input_names)
        _ = placeholders, inputs


    def build(self, *args, **kwargs):
        """ Build the model

        #. Create a graph
        #. Define ``is_training`` and ``global_step`` tensors
        #. Define a model architecture by calling :meth:``._build``
        #. Create a loss function from config
        #. Create an optimizer and define a train step from config
        #. `Set UPDATE_OPS control dependency on train step
           <https://www.tensorflow.org/api_docs/python/tf/layers/batch_normalization>`_
        #. Create a tensorflow session
        """
        with self.graph.as_default():
            self.store_to_attr('is_training', tf.placeholder(tf.bool, name='is_training'))
            self.store_to_attr('global_step', tf.Variable(0, trainable=False, name='global_step'))

            self._build()

            self._make_loss()
            self.store_to_attr('loss', tf.losses.get_total_loss())

            if self.train_step is None:
                optimizer = self._make_optimizer()

                update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
                with tf.control_dependencies(update_ops):
                    self.store_to_attr('train_step', optimizer.minimize(self.loss, global_step=self.global_step))
            else:
                self.store_to_attr('train_step', self.train_step)

            session_config = self.get_from_config('session', {})
            self.session = tf.Session(**session_config)
            self.session.run(tf.global_variables_initializer())

    def _make_inputs(self, names=None):
        """ Make model input data using config

        In the config's inputs section it looks for ``names``, creates placeholders required, and
        makes some typical transformations (like one-hot-encoding and reshaping), if needed.

        **Configuration**

        inputs : dict
            - key : str
                a placeholder name
            - values : dict or tuple
                each input's config

        Input config:

        ``dtype`` : str or tf.DType (by default 'float32')
            data type

        ``shape`` : int, tuple, list or None (default)
            a desired tensor shape which includes the number of channels/classes and doesn't include a batch size.

        ``classes`` : array-like or None (default)
            an array of class labels if data labels are strings or anything else except ``np.arange(num_classes)``

        ``data_format`` : str {``'channels_first'``, ``'channels_last'``} or {``'f'``, ``'l'``}
            The ordering of the dimensions in the inputs. Default is 'channels_last'.

        ``transform`` : str or callable
            Predefined transforms are

            - ``'ohe'`` - one-hot encoding
            - ``'mip @ d'`` - maximum intensity projection :func:`~.layers.mip`
              with depth ``d`` (should be int)

        ``name`` : str
            a name for the transformed and reshaped tensor.

        If an input config is a tuple, it should contain all items exactly in the order shown above:
        dtype, shape, classes, data_format, transform, name.
        If an item is None, the default value will be used instead.

        **How it works**

        Input data will be cast to ``dtype``, then reshaped to ``shape`` in accordance with ``data_format``
        and finally transformed with a ``transform`` function. The resulting tensor will have the name ``name``.
        The original placeholder tensor will have the name ``key``.

        Parameters
        ----------
        names : list
            placeholder names that are expected in the config inputs section

        Raises
        ------
        KeyError if there is any name missing in the config's input section.
        ValueError if there are duplicate names.

        Returns
        -------
        placeholders : dict
            key : str
                a placeholder name
            value : tf.Tensor
                placeholder tensor
        tensors : dict
            key : str
                a placeholder name
            value : tf.Tensor
                placeholder tensor after reshaping and transformations
        """
        config = self.get_from_config('inputs') or {}
        config = copy1(config)

        names = names or []
        missing_names = set(names) - set(config.keys())
        if len(missing_names) > 0:
            raise KeyError("Inputs should contain {} names".format(missing_names))

        placeholder_names = set(config.keys())
        tensor_names = set(x.get('name') for x in config.values() if x.get('name'))
        wrong_names = placeholder_names & tensor_names
        if len(wrong_names) > 0:
            raise ValueError('Inputs contain duplicate names:', wrong_names)

        param_names = ('dtype', 'shape', 'classes', 'data_format', 'transform', 'name')
        defaults = dict(data_format='channels_last')

        placeholders = dict()
        tensors = dict()

        for input_name, input_config in config.items():
            if isinstance(input_config, (tuple, list)):
                input_config = list(input_config) + [None for _ in param_names]
                input_config = input_config[:len(param_names)]
                input_config = dict(zip(param_names, input_config))
                input_config = dict((k, v) for k, v in input_config.items() if v is not None)
            input_config = {**defaults, **input_config}

            shape = input_config.get('shape')
            if isinstance(shape, int):
                input_config['shape'] = (shape,)

            self._inputs[input_name] = dict(config=input_config)

            if self.has_classes(input_name):
                dtype = input_config.get('dtype', tf.int32)
            else:
                dtype = input_config.get('dtype', 'float')
            tensor = tf.placeholder(dtype, name=input_name)
            placeholders[input_name] = tensor

            if input_config.get('data_format') == 'l':
                input_config['data_format'] = 'channels_last'
            elif input_config.get('data_format') == 'f':
                input_config['data_format'] = 'channels_first'

            self._inputs[input_name] = dict(config=input_config)
            tensor = self._make_transform(input_name, tensor, input_config)

            shape = input_config.get('shape')
            if input_config.get('__reshape', True) and isinstance(shape, (list, tuple)):
                input_config['shape'] = shape
                tensor = tf.reshape(tensor, [-1] + list(shape))

            name = input_config.get('name')
            if name is not None:
                tensor = tf.identity(tensor, name=name)

            tensors[input_name] = tensor

            self._inputs[input_name] = dict(config=input_config, placeholder=placeholders[input_name], tensor=tensor)

        return placeholders, tensors

    def _make_transform(self, input_name, tensor, config):
        if config is not None:
            transform_name = config.get('transform')
            if isinstance(transform_name, str):
                transforms = {
                    'ohe': self._make_ohe,
                    'mip': self._make_mip
                }

                if transform_name.startswith('mip'):
                    parts = transform_name.split('@')
                    transform_name = parts[0].strip()
                    config['depth'] = int(parts[1])
                    # mip has to know shape, so we need call reshape first
                    shape = config.get('shape')
                    tensor = tf.reshape(tensor, [-1] + list(shape))
                    # do not make reshape twice
                    config.update({'__reshape': False})

                tensor = transforms[transform_name](input_name, tensor, config)
            elif callable(transform_name):
                tensor = transform_name(tensor)
            elif transform_name is not None:
                raise ValueError("Unknown transform {}".format(transform_name))
        return tensor

    def _make_ohe(self, input_name, tensor, config):
        if config.get('shape') is None and config.get('classes') is None:
            raise ValueError("shape and classes cannot be both None for input '{}'".format(input_name))

        num_classes = self.num_classes(input_name)
        axis = -1 if self.data_format(input_name) else 1
        tensor = tf.one_hot(tensor, depth=num_classes, axis=axis)
        return tensor

    def to_classes(self, tensor, input_name, name=None):
        """ Convert tensor with labels to classes of ``input_name`` """
        if tensor.dtype in [tf.float16, tf.float32, tf.float64]:
            tensor = tf.argmax(tensor, axis=-1, name=name)
        if self.has_classes(input_name):
            self._to_classes.update({tensor: input_name})
        return tensor

    def _make_mip(self, input_name, tensor, config):
        depth = config.get('depth')
        if depth is None:
            raise ValueError("mip should be specifies as mip @ depth")
        tensor = mip(tensor, depth=depth, data_format=self.data_format(input_name))
        return tensor

    def _unpack_fn_from_config(self, param, default=None):
        par = self.get_from_config(param, default)

        if par is None:
            return None, None

        if isinstance(par, (tuple, list)):
            if len(par) == 0:
                par_name = None
            elif len(par) == 1:
                par_name, par_args = par[0], {}
            elif len(par) == 2:
                par_name, par_args = par
            else:
                par_name, par_args = par[0], par[1:]
        elif isinstance(par, dict):
            par = par.copy()
            par_name, par_args = par.pop('name', None), par
        else:
            par_name, par_args = par, {}

        return par_name, par_args

    def _make_loss(self):
        """ Return a loss function from config """
        loss = self.get_from_config("loss")
        if loss is None:
            if len(tf.losses.get_losses()) == 0:
                raise ValueError("Loss is not defined in the model %s" % self)
        elif isinstance(loss, str) and hasattr(tf.losses, loss):
            loss = getattr(tf.losses, loss)
        elif isinstance(loss, str):
            loss = LOSSES.get(re.sub('[-_ ]', '', loss).lower(), None)
        elif callable(loss):
            pass
        else:
            raise ValueError("Unknown loss", loss)

        if loss is not None:
            try:
                predictions = self.graph.get_tensor_by_name("predictions:0")
                targets = self.graph.get_tensor_by_name("targets:0")
            except KeyError:
                raise KeyError("Model %s does not have 'predictions' or 'targets' tensors" % self.name)
            else:
                tf.losses.add_loss(loss(targets, predictions))

    def _make_decay(self):
        decay_name, decay_args = self._unpack_fn_from_config('decay')

        if decay_name is None:
            pass
        elif callable(decay_name):
            pass
        elif isinstance(decay_name, str) and hasattr(tf.train, decay_name):
            decay_name = getattr(tf.train, decay_name)
        elif decay_name in DECAYS:
            decay_name = DECAYS.get(re.sub('[-_ ]', '', decay_name).lower(), None)
        else:
            raise ValueError("Unknown learning rate decay method", decay_name)

        return decay_name, decay_args

    def _make_optimizer(self):
        optimizer_name, optimizer_args = self._unpack_fn_from_config('optimizer', 'Adam')

        if callable(optimizer_name):
            pass
        elif isinstance(optimizer_name, str) and hasattr(tf.train, optimizer_name):
            optimizer_name = getattr(tf.train, optimizer_name)
        elif isinstance(optimizer_name, str) and hasattr(tf.train, optimizer_name + 'Optimizer'):
            optimizer_name = getattr(tf.train, optimizer_name + 'Optimizer')
        else:
            raise ValueError("Unknown optimizer", optimizer_name)

        decay_name, decay_args = self._make_decay()
        if decay_name is not None:
            optimizer_args['learning_rate'] = decay_name(**decay_args, global_step=self.global_step)

        optimizer = optimizer_name(**optimizer_args)

        return optimizer

    def get_number_of_trainable_vars(self):
        """ Return the number of trainable variable in the model graph """
        with self.graph:
            arr = np.asarray([np.prod(self.get_shape(v)) for v in tf.trainable_variables()])
        return np.sum(arr)

    def get_tensor_config(self, tensor, **kwargs):
        """ Return tensor configuration

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        dict
            tensor config (see :meth:`._make_inputs`)

        Raises
        ------
        ValueError shape in tensor configuration isn't int, tuple or list
        """
        if isinstance(tensor, tf.Tensor):
            names = [n for n, i in self._inputs.items() if tensor in [i['placeholder'], i['tensor']]]
            if len(names) > 0:
                input_name = names[0]
            else:
                input_name = tensor.name
        elif isinstance(tensor, str):
            if tensor in self._inputs:
                input_name = tensor
            else:
                input_name = self._map_name(tensor)
        else:
            raise TypeError("Tensor can be tf.Tensor or string, but given %s" % type(tensor))

        if input_name in self._inputs:
            config = self._inputs[input_name]['config']
            shape = config.get('shape')
            if isinstance(shape, int):
                shape = (shape,)
            kwargs['shape'] = shape

            if 'classes' not in config and 'shape' not in config:
                raise ValueError("one of 'classes' or 'shape' must be present in inputs config")
            if 'classes' not in config:
                if not isinstance(shape, (list, tuple)):
                    raise ValueError('shape must be int, tuple or list but {} was given'.format(type(shape)))
        else:
            tensor = self.graph.get_tensor_by_name(input_name)
            shape = tensor.get_shape().as_list()[1:]
            config = dict(dtype=tensor.dtype, shape=shape, name=tensor.name, data_format='channels_last')
            config = {**config, **kwargs}

        return config


    def num_channels(self, tensor, **kwargs):
        """ Return the number of channels in the tensor

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        number of channels : int
        """
        config = self.get_tensor_config(tensor, **kwargs)
        shape = config.get('shape')
        data_format = config.get('data_format')
        channels_dim = -1 if data_format == "channels_last" or not data_format.startswith("NC") else 0
        return shape[channels_dim]

    def has_classes(self, tensor):
        """ Check if a tensor has classes defined in the config """
        config = self.get_tensor_config(tensor)
        has = config.get('classes') is not None
        return has

    def classes(self, tensor):
        """ Return the  number of classes """
        config = self.get_tensor_config(tensor)
        return config.get('classes')

    def num_classes(self, tensor):
        """ Return the  number of classes """
        if self.has_classes(tensor):
            return len(self.classes(tensor))
        return self.num_channels(tensor)

    def spatial_dim(self, tensor, **kwargs):
        """ Return the tensor spatial dimensionality (without channels dimension)

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        number of spatial dimensions : int
        """
        config = self.get_tensor_config(tensor, **kwargs)
        return len(config.get('shape')) - 1

    def data_format(self, tensor, **kwargs):
        """ Return the tensor data format (channels_last or channels_first)

        Parameters
        ----------
        tensor : str or tf.Tensor

        Returns
        -------
        data_format : str
        """
        config = self.get_tensor_config(tensor, **kwargs)
        return config.get('data_format')
{'fraction_non_alphanumeric': 0.11288805268109126, 'fraction_numerical': 0.002266313178824938, 'mean_word_length': 2.154582855601268, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 27, 'https://': 18, 'lorem ipsum': 0, 'www.': 18, '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': '2940793', 'n_tokens_mistral': 13245, 'n_tokens_neox': 12635, 'n_words': 3703}