<commit_before># coding=utf-8
# Copyright 2017 The Tensor2Tensor 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.

"""Library for training on TPU. See tpu_trainer.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

from tensor2tensor.utils import t2t_model
from tensor2tensor.utils import trainer_utils

import tensorflow as tf


def create_run_config(master="",
                      model_dir=None,
                      iterations_per_loop=1000,
                      num_shards=8,
                      log_device_placement=False,
                      save_checkpoints_steps=1000,
                      num_gpus=1,
                      gpu_order="",
                      shard_to_cpu=False,
                      num_async_replicas=1,
                      use_tpu=True):
  """Create TPUConfig and tpu.RunConfig."""
  session_config = tf.ConfigProto(
      allow_soft_placement=True, log_device_placement=log_device_placement)
  run_config_args = {
      "model_dir": model_dir,
      "session_config": session_config,
      "save_summary_steps": 0,
      "save_checkpoints_steps": save_checkpoints_steps,
  }
  run_config_cls = tf.estimator.RunConfig

  # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args
  if use_tpu:
    run_config_cls = tf.contrib.tpu.RunConfig
    tpu_config = tf.contrib.tpu.TPUConfig(
        iterations_per_loop=iterations_per_loop,
        num_shards=num_shards,
        per_host_input_for_training=(num_shards <= 8))
    run_config_args["master"] = master
    run_config_args["tpu_config"] = tpu_config

  config = run_config_cls(**run_config_args)

  # If not using TPU, add device info for data_parallelism
  config.use_tpu = use_tpu
  if not use_tpu:
    config.t2t_device_info = {
        "num_gpus": num_gpus,
        "gpu_order": gpu_order,
        "shard_to_cpu": shard_to_cpu,
        "num_shards": max(1, num_gpus + int(shard_to_cpu)),
        "num_async_replicas": num_async_replicas,
    }

  return config


def create_estimator(model_name,
                     hparams,
                     run_config,
                     schedule="train_and_evaluate",
                     use_tpu=True):
  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
      model_name, hparams, use_tpu=use_tpu)

  if use_tpu:
    batch_size = hparams.tpu_batch_size_per_shard
    batch_size *= run_config.tpu_config.num_shards
    eval_batch_size = batch_size * 2
    if schedule == "train":
      # Estimator takes the presence of eval_batch_size as an indication that
      # an eval is being performed, and complains about num_shards being too
      # big. So we have to set eval_batch_size to None.
      eval_batch_size = None
    return tf.contrib.tpu.TPUEstimator(
        model_fn=model_fn,<commit_msg>Add empty t2t_device_info to config in TPU path<commit_after># coding=utf-8
# Copyright 2017 The Tensor2Tensor 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.

"""Library for training on TPU. See tpu_trainer.py."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports

from tensor2tensor.utils import t2t_model
from tensor2tensor.utils import trainer_utils

import tensorflow as tf


def create_run_config(master="",
                      model_dir=None,
                      iterations_per_loop=1000,
                      num_shards=8,
                      log_device_placement=False,
                      save_checkpoints_steps=1000,
                      num_gpus=1,
                      gpu_order="",
                      shard_to_cpu=False,
                      num_async_replicas=1,
                      use_tpu=True):
  """Create TPUConfig and tpu.RunConfig."""
  session_config = tf.ConfigProto(
      allow_soft_placement=True, log_device_placement=log_device_placement)
  run_config_args = {
      "model_dir": model_dir,
      "session_config": session_config,
      "save_summary_steps": 0,
      "save_checkpoints_steps": save_checkpoints_steps,
  }
  run_config_cls = tf.estimator.RunConfig

  # If using TPU, use TPU RunConfig, add TPUConfig, and add additional args
  if use_tpu:
    run_config_cls = tf.contrib.tpu.RunConfig
    tpu_config = tf.contrib.tpu.TPUConfig(
        iterations_per_loop=iterations_per_loop,
        num_shards=num_shards,
        per_host_input_for_training=(num_shards <= 8))
    run_config_args["master"] = master
    run_config_args["tpu_config"] = tpu_config

  config = run_config_cls(**run_config_args)

  # If not using TPU, add device info for data_parallelism
  config.use_tpu = use_tpu
  config.t2t_device_info = {}
  if not use_tpu:
    config.t2t_device_info = {
        "num_gpus": num_gpus,
        "gpu_order": gpu_order,
        "shard_to_cpu": shard_to_cpu,
        "num_shards": max(1, num_gpus + int(shard_to_cpu)),
        "num_async_replicas": num_async_replicas,
    }

  return config


def create_estimator(model_name,
                     hparams,
                     run_config,
                     schedule="train_and_evaluate",
                     use_tpu=True):
  model_fn = t2t_model.T2TModel.make_estimator_model_fn(
      model_name, hparams, use_tpu=use_tpu)

  if use_tpu:
    batch_size = hparams.tpu_batch_size_per_shard
    batch_size *= run_config.tpu_config.num_shards
    eval_batch_size = batch_size * 2
    if schedule == "train":
      # Estimator takes the presence of eval_batch_size as an indication that
      # an eval is being performed, and complains about num_shards being too
      # big. So we have to set eval_batch_size to None.
      eval_batch_size = None
    return tf.contrib.tpu.TPUEstimator(
        model_fn=model_fn,
{'fraction_non_alphanumeric': 0.07174490699734277, 'fraction_numerical': 0.0094478889873044, 'mean_word_length': 2.8494318181818183, 'pattern_counts': {'":': 20, '<': 5, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4135633', 'n_tokens_mistral': 2156, 'n_tokens_neox': 2073, 'n_words': 535}