<commit_before>from __future__ import absolute_import, division, print_function

import copy
import inspect
import os
import shutil
from tempfile import mkdtemp
from types import ModuleType
from typing import Callable

import dill
from six import string_types
from tensorflow.python import saved_model
from tensorflow.python.keras import Model, Sequential
from tensorflow.python.keras import layers as layer_module
from tensorflow.python.keras.engine import base_layer_utils
from tensorflow.python.keras.layers import Activation, Dense, Layer
from tensorflow.python.keras.utils import tf_utils
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import keras_export

from odin.utils import get_all_files

__all__ = ['AdvanceModel', 'Identity', 'Parallel']


class AdvanceModel(Model):
  """ The advance model improving the serialization and deserialization of
  complex Model

  Parameters
  ----------
  parameters : {`dict` or `None`}
    recording the arguments given to `__init__`, recommend passing
    the `locals()` dictionary
  """

  def __init__(self, parameters=None, name=None):
    super(AdvanceModel, self).__init__(name=name)
    self.supports_masking = True
    self._build_input_shape = None

    if parameters is None:
      parameters = {}
    else:
      parameters = dict(parameters)
    kwargs = parameters.pop('kwargs', {})
    parameters.update(kwargs)
    parameters.pop('self', None)
    parameters.pop('__class__', None)
    # tricky recursive reference of overriding classes
    parameters.pop('parameters', None)
    self._parameters = parameters
    self._optimizer_weights = None

  @property
  def optimizer(self):
    if self._optimizer_weights is not None and \
      self._optimizer is not None and \
        len(self._optimizer.weights) == len(self._optimizer_weights):
      # TODO: this is bad solution to resume the training
      self._optimizer.set_weights(self._optimizer_weights)
      self._optimizer_weights = None
    return self._optimizer

  @optimizer.setter
  def optimizer(self, optz):
    self._optimizer = optz

  def __getstate__(self):
    configs = dill.dumps(self.get_config())
    weights = dill.dumps(self.get_weights())
    optimizer = self.optimizer
    if optimizer is not None:
      optimizer = [
          optimizer.__class__,
          optimizer.get_config(),
          optimizer.get_weights(),
      ]
      optimizer = dill.dumps(optimizer)
    return configs, weights, optimizer

  def __setstate__(self, states):
    configs, weights, optimizer = states
    configs = dill.loads(configs)
    weights = dill.loads(weights)
    clone = self.__class__.from_config(configs)
    self.__dict__.update(clone.__dict__)
    self.set_weights(weights)
    if optimizer is not None:
      cls, configs, weights = dill.loads(optimizer)
      optimizer = cls.from_config(configs)
      self.compile(optimizer)
      self._optimizer_weights = weights

  @property
  def custom_objects(self):
    """ This property could be overrided to provide custom Layer class
    Expect a dictionary mapping from class name to the class itself for
    deserialization, or, list of modules
    """
    return {}

  @property
  def parameters(self):
    return dict(self._parameters)

  @base_layer_utils.default
  def build(self, input_shape=None):
    # if default is enable, it mean the build method is not overrided
    # and won't be called during __call__
    if self._is_graph_network:
      self._init_graph_network(self.inputs, self.outputs, name=self.name)
    else:
      if input_shape is None:
        if self._build_input_shape is None:
          raise ValueError('You must provide an `input_shape` argument.')
        else:
          input_shape = self._build_input_shape
      input_shape = tuple(input_shape)
      self._build_input_shape = input_shape
      super(AdvanceModel, self).build(input_shape)
    self.built = True

  def __call__(self, inputs, *args, **kwargs):
    # This to make sure input_shapes is recorded after every call
    if self._build_input_shape is None:
      input_list = nest.flatten(inputs)
      input_shapes = None
      if all(hasattr(x, 'shape') for x in input_list):
        input_shapes = nest.map_structure(lambda x: x.shape, inputs)
      self._build_input_shape = input_shapes
    return super(AdvanceModel, self).__call__(inputs, *args, **kwargs)

  def get_config(self):
    source = inspect.getsource(self.__class__)
    default_keys = []
    for t in type.mro(type(self)):
      default_keys += dir(t)

    attributes = {}
    layer_attributes = {}
    for key, val in self.__dict__.items():
      if isinstance(val, Layer):
        layer_attributes[id(val)] = key
      elif 'self.' + key in source and key not in default_keys:
        try:
          attr = getattr(self, key)
          if not inspect.ismethod(attr) and \
            not isinstance(attr, property) and \<commit_msg>improve attributes tracking for AdvancedModel<commit_after>from __future__ import absolute_import, division, print_function

import copy
import inspect
import os
import shutil
from tempfile import mkdtemp
from types import ModuleType
from typing import Callable

import dill
import tensorflow as tf
from six import string_types
from tensorflow.python import saved_model
from tensorflow.python.keras import Model, Sequential
from tensorflow.python.keras import layers as layer_module
from tensorflow.python.keras.engine import base_layer_utils
from tensorflow.python.keras.layers import Activation, Dense, Layer
from tensorflow.python.keras.utils import tf_utils
from tensorflow.python.util import nest
from tensorflow.python.util.tf_export import keras_export

from odin.utils import get_all_files

__all__ = ['AdvanceModel', 'Identity', 'Parallel']


class AdvanceModel(Model):
  """ The advance model improving the serialization and deserialization of
  complex Model

  Parameters
  ----------
  parameters : {`dict` or `None`}
    recording the arguments given to `__init__`, recommend passing
    the `locals()` dictionary
  """

  def __init__(self, parameters=None, name=None):
    super(AdvanceModel, self).__init__(name=name)
    self.supports_masking = True
    self._build_input_shape = None

    if parameters is None:
      parameters = {}
    else:
      parameters = dict(parameters)
    kwargs = parameters.pop('kwargs', {})
    parameters.update(kwargs)
    parameters.pop('self', None)
    parameters.pop('__class__', None)
    # tricky recursive reference of overriding classes
    parameters.pop('parameters', None)
    self._parameters = parameters
    self._optimizer_weights = None
    self._optimizer = None

  @property
  def optimizer(self):
    if self._optimizer_weights is not None and \
      self._optimizer is not None and \
        len(self._optimizer.weights) == len(self._optimizer_weights):
      # TODO: this is bad solution to resume the training
      self._optimizer.set_weights(self._optimizer_weights)
      self._optimizer_weights = None
    return self._optimizer

  @optimizer.setter
  def optimizer(self, optz):
    self._optimizer = optz

  def __getstate__(self):
    configs = dill.dumps(self.get_config())
    weights = dill.dumps(self.get_weights())
    optimizer = self.optimizer
    if optimizer is not None:
      optimizer = [
          optimizer.__class__,
          optimizer.get_config(),
          optimizer.get_weights()
          if self._optimizer_weights is None else self._optimizer_weights,
      ]
      optimizer = dill.dumps(optimizer)
    return configs, weights, optimizer

  def __setstate__(self, states):
    configs, weights, optimizer = states
    configs = dill.loads(configs)
    weights = dill.loads(weights)
    clone = self.__class__.from_config(configs)
    self.__dict__.update(clone.__dict__)
    self.set_weights(weights)
    if optimizer is not None:
      cls, configs, weights = dill.loads(optimizer)
      optimizer = cls.from_config(configs)
      self.compile(optimizer)
      self._optimizer_weights = weights

  @property
  def custom_objects(self):
    """ This property could be overrided to provide custom Layer class
    Expect a dictionary mapping from class name to the class itself for
    deserialization, or, list of modules
    """
    return {}

  @property
  def parameters(self):
    return dict(self._parameters)

  @base_layer_utils.default
  def build(self, input_shape=None):
    # if default is enable, it mean the build method is not overrided
    # and won't be called during __call__
    if self._is_graph_network:
      self._init_graph_network(self.inputs, self.outputs, name=self.name)
    else:
      if input_shape is None:
        if self._build_input_shape is None:
          raise ValueError('You must provide an `input_shape` argument.')
        else:
          input_shape = self._build_input_shape
      # do not convert input_shape to tuple, multiple inputs
      # will create a ListWrapper of multiple TensorShape.
      self._build_input_shape = input_shape
      super(AdvanceModel, self).build(input_shape)
    self.built = True

  def __call__(self, inputs, *args, **kwargs):
    # This to make sure input_shapes is recorded after every call
    if self._build_input_shape is None:
      input_list = nest.flatten(inputs)
      input_shapes = None
      if all(hasattr(x, 'shape') for x in input_list):
        input_shapes = nest.map_structure(lambda x: x.shape, inputs)
      self._build_input_shape = input_shapes
    return super(AdvanceModel, self).__call__(inputs, *args, **kwargs)

  def get_config(self):
    subclass_types = []
    for t in type.mro(type(self)):
      if t == AdvanceModel:
        break
      subclass_types.append(t)
    default_types = [t for t in type.mro(type(self)) if t not in subclass_types]

    default_keys = []
    default_source = ''
    for t in default_types:
      default_keys += dir(t)
      # some class is builtin, so impossible to get the source code
      try:
        default_source += inspect.getsource(t)
      except:
        pass

    source = ''
    for t in subclass_types:
      source += inspect.getsource(t)

    attributes = {}
    layer_attributes = {}  # mapping attribute_name -> Layer
    for key, val in self.__dict__.items():
      # TODO: not a good solution here, history make recursive reference
      # but couldn't found in Model soruce code where it is setted
      if key in ('_parameters', 'history') or key in self.parameters:
        continue
      if isinstance(val, Layer):
        layer_attributes[id(val)] = key
      elif isinstance(val, tf.Variable):
        pass
      elif key not in default_keys and \
        'self.' + key not in default_source and \
          'self.' + key in source:
        try:
          attr = getattr(self, key)
          if not inspect.ismethod(attr) and \
            not isinstance(attr, property) and \
{'fraction_non_alphanumeric': 0.08485351022664456, 'fraction_numerical': 0.0, 'mean_word_length': 3.282051282051282, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, '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': '3489182', 'n_tokens_mistral': 3155, 'n_tokens_neox': 3105, 'n_words': 932}