<commit_before># coding=utf-8
# Copyright 2019 The Tensor2Robot Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""TFModel abstract subclasses."""

from __future__ import absolute_import
from __future__ import division

from __future__ import print_function

import abc
import os
from absl import flags
from absl import logging
import gin
from gin.tf import utils as gin_utils
import numpy as np

from tensor2robot.models import model_interface
from tensor2robot.models import optimizers
from tensor2robot.preprocessors import abstract_preprocessor
from tensor2robot.preprocessors import noop_preprocessor
from tensor2robot.utils import tensorspec_utils
import tensorflow as tf

from typing import Any, Callable, Dict, Optional, Text, Tuple, Union

FLAGS = flags.FLAGS
TRAIN = tf.estimator.ModeKeys.TRAIN
EVAL = tf.estimator.ModeKeys.EVAL
PREDICT = tf.estimator.ModeKeys.PREDICT

RunConfigType = Optional[
    Union[tf.estimator.RunConfig, tf.contrib.tpu.RunConfig]]
ParamsType = Optional[Dict[Text, Any]]

DictOrSpec = Union[Dict[Text, tf.Tensor], tensorspec_utils.TensorSpecStruct]

ModelTrainOutputType = Union[tf.Tensor, Tuple[tf.Tensor, DictOrSpec]]
ExportOutputType = Union[Dict[Text, tf.Tensor], Tuple[
    Dict[Text, tf.Tensor], Dict[Text, tf.estimator.export.PredictOutput]]]

try:
  flags.DEFINE_string('master', '', 'Master for TPU RunConfig')
except flags.DuplicateFlagError:
  pass

DEVICE_TYPE_CPU = 'cpu'
DEVICE_TYPE_GPU = 'gpu'
DEVICE_TYPE_TPU = 'tpu'

gin_configurable_run_config_cls = gin.external_configurable(
    tf.estimator.RunConfig,
    name='tf.estimator.RunConfig',
    blacklist=['model_dir'])

gin_configurable_tpu_run_config_cls = gin.external_configurable(
    tf.contrib.tpu.RunConfig,
    name='tf.contrib.tpu.RunConfig',
    blacklist=['model_dir', 'tpu_config'])

gin_configurable_tpu_config_cls = gin.external_configurable(
    tf.contrib.tpu.TPUConfig, name='tf.contrib.tpu.TPUConfig')


@gin.configurable
def default_init_from_checkpoint_fn(checkpoint):
  """init_from_checkpoint_fn that can be used to init a model from a checkpoint.

  Args:
    checkpoint: String pointing to path of TF checkpoint.
  """
  logging.info('Initializing model weights from %s', checkpoint)
  variables_to_restore = tf.contrib.framework.get_trainable_variables()
  assignment_map = {v.op.name: v for v in variables_to_restore}
  tf.train.init_from_checkpoint(checkpoint, assignment_map)


@gin.configurable
class AbstractT2RModel(model_interface.ModelInterface):<commit_msg>Support initializing from partial checkpoints, by default this is disabled.<commit_after># coding=utf-8
# Copyright 2019 The Tensor2Robot Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""TFModel abstract subclasses."""

from __future__ import absolute_import
from __future__ import division

from __future__ import print_function

import abc
import os
from absl import flags
from absl import logging
import gin
from gin.tf import utils as gin_utils
import numpy as np

from tensor2robot.models import model_interface
from tensor2robot.models import optimizers
from tensor2robot.preprocessors import abstract_preprocessor
from tensor2robot.preprocessors import noop_preprocessor
from tensor2robot.utils import tensorspec_utils
import tensorflow as tf

from typing import Any, Callable, Dict, Optional, Text, Tuple, Union

FLAGS = flags.FLAGS
TRAIN = tf.estimator.ModeKeys.TRAIN
EVAL = tf.estimator.ModeKeys.EVAL
PREDICT = tf.estimator.ModeKeys.PREDICT

RunConfigType = Optional[
    Union[tf.estimator.RunConfig, tf.contrib.tpu.RunConfig]]
ParamsType = Optional[Dict[Text, Any]]

DictOrSpec = Union[Dict[Text, tf.Tensor], tensorspec_utils.TensorSpecStruct]

ModelTrainOutputType = Union[tf.Tensor, Tuple[tf.Tensor, DictOrSpec]]
ExportOutputType = Union[Dict[Text, tf.Tensor], Tuple[
    Dict[Text, tf.Tensor], Dict[Text, tf.estimator.export.PredictOutput]]]

try:
  flags.DEFINE_string('master', '', 'Master for TPU RunConfig')
except flags.DuplicateFlagError:
  pass

DEVICE_TYPE_CPU = 'cpu'
DEVICE_TYPE_GPU = 'gpu'
DEVICE_TYPE_TPU = 'tpu'

gin_configurable_run_config_cls = gin.external_configurable(
    tf.estimator.RunConfig,
    name='tf.estimator.RunConfig',
    blacklist=['model_dir'])

gin_configurable_tpu_run_config_cls = gin.external_configurable(
    tf.contrib.tpu.RunConfig,
    name='tf.contrib.tpu.RunConfig',
    blacklist=['model_dir', 'tpu_config'])

gin_configurable_tpu_config_cls = gin.external_configurable(
    tf.contrib.tpu.TPUConfig, name='tf.contrib.tpu.TPUConfig')


@gin.configurable
def default_init_from_checkpoint_fn(checkpoint,
                                    allow_partial_restore = False):
  """init_from_checkpoint_fn that can be used to init a model from a checkpoint.

  Args:
    checkpoint: String pointing to path of TF checkpoint.
    allow_partial_restore: If True, we allow partial restore, otherwise we raise
      an error if a variable cannot be restored.

  Raises:
    A ValueError if a variable(s) is missing and partial restore is not
    explicitly enabled.
  """
  logging.info('Initializing model weights from %s', checkpoint)
  reader = tf.train.load_checkpoint(checkpoint)
  variables_to_restore = tf.contrib.framework.get_trainable_variables()
  assignment_map = {}
  for v in variables_to_restore:
    op_name = v.op.name
    if reader.has_tensor(op_name):
      logging.info('Loading variable %s from checkpoint', op_name)
      assignment_map[op_name] = v
    elif allow_partial_restore:
      logging.warning('Variable %s is not in the checkpoint, skipping.',
                      op_name)
    else:
      raise ValueError('Attempting to restore variable {} which is '
                       'not in the checkpoint.'.format(op_name))

  tf.train.init_from_checkpoint(checkpoint, assignment_map)


@gin.configurable
class AbstractT2RModel(model_interface.ModelInterface):
{'fraction_non_alphanumeric': 0.08277372262773723, 'fraction_numerical': 0.004671532846715329, 'mean_word_length': 4.931601731601732, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '628611', 'n_tokens_mistral': 2104, 'n_tokens_neox': 2060, 'n_words': 605}