# Copyright 2020, Google LLC.
#
# 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.
"""Runs federated training with differential privacy on various tasks."""

import functools

from absl import app
from absl import flags
from absl import logging
import tensorflow as tf
import tensorflow_federated as tff

from utils import task_utils
from utils import training_utils
from utils import utils_impl
from utils.optimizers import optimizer_utils

with utils_impl.record_hparam_flags() as optimizer_flags:
  # Defining optimizer flags
  optimizer_utils.define_optimizer_flags('client')
  optimizer_utils.define_optimizer_flags('server')

with utils_impl.record_hparam_flags() as shared_flags:
  # Federated training hyperparameters
  flags.DEFINE_integer('client_epochs_per_round', 1,
                       'Number of epochs in the client to take per round.')
  flags.DEFINE_integer('client_batch_size', 20, 'Batch size on the clients.')
  flags.DEFINE_integer('clients_per_round', 10,
                       'How many clients to sample per round.')
  flags.DEFINE_integer('client_datasets_random_seed', 1,
                       'Random seed for client sampling.')
  flags.DEFINE_integer(
      'max_elements_per_client', None, 'Maximum number of '
      'elements for each training client. If set to None, all '
      'available examples are used.')

  # Training loop configuration
  flags.DEFINE_integer('total_rounds', 200, 'Number of total training rounds.')
  flags.DEFINE_string(
      'experiment_name', None, 'The name of this experiment. Will be append to '
      '--root_output_dir to separate experiment results.')
  flags.DEFINE_string('root_output_dir', '/tmp/fed_opt/',
                      'Root directory for writing experiment output.')
  flags.DEFINE_integer(
      'rounds_per_eval', 1,
      'How often to evaluate the global model on the validation dataset.')
  flags.DEFINE_integer(
      'num_validation_examples', -1, 'The number of validation'
      'examples to use. If set to -1, all available examples '
      'are used.')
  flags.DEFINE_integer('rounds_per_checkpoint', 50,
                       'How often to checkpoint the global model.')

with utils_impl.record_hparam_flags() as dp_flags:
  # Differential privacy flags
  flags.DEFINE_float(
      'clip', None, 'Clip value for fixed clipping or initial clip for '
      'adaptive clipping. If None, no clipping is used.')
  flags.DEFINE_float('noise_multiplier', None,
                     'Noise multiplier. If None, non-DP aggregator is used.')
  flags.DEFINE_float(
      'adaptive_clip_learning_rate', None, 'Adaptive clip learning rate. If '
      'None, clip adaptation is not used.')
  flags.DEFINE_float('target_unclipped_quantile', 0.5,
                     'Target unclipped quantile.')
  flags.DEFINE_boolean('uniform_weighting', False,
                       'Whether to weigh clients uniformly.')

# Task specification
with utils_impl.record_hparam_flags() as task_flags:
  task_utils.define_task_flags()

FLAGS = flags.FLAGS


def _write_hparam_flags():
  """Returns an ordered dictionary of pertinent hyperparameter flags."""
  hparam_dict = utils_impl.lookup_flag_values(shared_flags)

  # Update with optimizer flags corresponding to the chosen optimizers.
  opt_flag_dict = utils_impl.lookup_flag_values(optimizer_flags)
  opt_flag_dict = optimizer_utils.remove_unused_flags('client', opt_flag_dict)
  opt_flag_dict = optimizer_utils.remove_unused_flags('server', opt_flag_dict)
  hparam_dict.update(opt_flag_dict)

  # Update with task flags
  task_flag_dict = utils_impl.lookup_flag_values(task_flags)
  hparam_dict.update(task_flag_dict)
  training_utils.write_hparams_to_csv(hparam_dict, FLAGS.root_output_dir,
                                      FLAGS.experiment_name)


def main(argv):
  if len(argv) > 1:
    raise app.UsageError('Expected no command-line arguments, '
                         'got: {}'.format(argv))

  client_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('client')
  server_optimizer_fn = optimizer_utils.create_optimizer_fn_from_flags('server')

  train_client_spec = tff.simulation.baselines.ClientSpec(
      num_epochs=FLAGS.client_epochs_per_round,
      batch_size=FLAGS.client_batch_size,
      max_elements=FLAGS.max_elements_per_client)
  task = task_utils.create_task_from_flags(train_client_spec)

  logging.info('Trainable weights:')
  for weight in task.model_fn().trainable_variables:
    logging.info('name: %s  shape: %s', weight.name, weight.shape)

  if FLAGS.uniform_weighting:
    client_weighting = tff.learning.ClientWeighting.UNIFORM
  elif FLAGS.task == 'shakespeare_character' or FLAGS.task == 'stackoverflow_word':

    def client_weighting(local_outputs):
      return tf.cast(tf.squeeze(local_outputs['num_tokens']), tf.float32)
  else:
    client_weighting = None

  if FLAGS.noise_multiplier is None:
    if FLAGS.uniform_weighting:
      aggregation_factory = tff.aggregators.UnweightedMeanFactory()
    else:
      aggregation_factory = tff.aggregators.MeanFactory()
    if FLAGS.clip is not None:
      if FLAGS.clip <= 0:
        raise ValueError('clip must be positive if clipping is enabled.')
      if FLAGS.adaptive_clip_learning_rate is None:
        clip = FLAGS.clip
      else:
        if FLAGS.adaptive_clip_learning_rate <= 0:
          raise ValueError('adaptive_clip_learning_rate must be positive if '
                           'adaptive clipping is enabled.')
        clip = tff.aggregators.PrivateQuantileEstimationProcess.no_noise(
            initial_estimate=FLAGS.clip,
            target_quantile=FLAGS.target_unclipped_quantile,
            learning_rate=FLAGS.adaptive_clip_learning_rate)
      aggregation_factory = tff.aggregators.clipping_factory(
          clip, aggregation_factory)
  else:
    if not FLAGS.uniform_weighting:
      raise ValueError(
          'Differential privacy is only implemented for uniform weighting.')
    if FLAGS.noise_multiplier <= 0:
      raise ValueError('noise_multiplier must be positive if DP is enabled.')
    if FLAGS.clip is None or FLAGS.clip <= 0:
      raise ValueError('clip must be positive if DP is enabled.')
    if FLAGS.adaptive_clip_learning_rate is None:
      aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_fixed(
          noise_multiplier=FLAGS.noise_multiplier,
          clients_per_round=FLAGS.clients_per_round,
          clip=FLAGS.clip)
    else:
      if FLAGS.adaptive_clip_learning_rate <= 0:
        raise ValueError('adaptive_clip_learning_rate must be positive if '
                         'adaptive clipping is enabled.')
      aggregation_factory = tff.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive(
          noise_multiplier=FLAGS.noise_multiplier,
          clients_per_round=FLAGS.clients_per_round,
          initial_l2_norm_clip=FLAGS.clip,
          target_unclipped_quantile=FLAGS.target_unclipped_quantile,
          learning_rate=FLAGS.adaptive_clip_learning_rate)

  iterative_process = tff.learning.build_federated_averaging_process(
      model_fn=task.model_fn,
      server_optimizer_fn=server_optimizer_fn,
      client_weighting=client_weighting,
      client_optimizer_fn=client_optimizer_fn,
      model_update_aggregation_factory=aggregation_factory)
  train_data = task.datasets.train_data.preprocess(
      task.datasets.train_preprocess_fn)
  training_process = (
      tff.simulation.compose_dataset_computation_with_iterative_process(
          train_data.dataset_computation, iterative_process))

  training_selection_fn = functools.partial(
      tff.simulation.build_uniform_sampling_fn(
          train_data.client_ids, random_seed=FLAGS.client_datasets_random_seed),
      size=FLAGS.clients_per_round)

  test_data = task.datasets.get_centralized_test_data()
  validation_data = test_data.take(FLAGS.num_validation_examples)
  federated_eval = tff.learning.build_federated_evaluation(task.model_fn)
  evaluation_selection_fn = lambda round_num: [validation_data]

  def evaluation_fn(state, evaluation_data):
    return federated_eval(state.model, evaluation_data)

  program_state_manager, metrics_managers = training_utils.create_managers(
      FLAGS.root_output_dir, FLAGS.experiment_name)
  _write_hparam_flags()
  state = tff.simulation.run_training_process(
      training_process=training_process,
      training_selection_fn=training_selection_fn,
      total_rounds=FLAGS.total_rounds,
      evaluation_fn=evaluation_fn,
      evaluation_selection_fn=evaluation_selection_fn,
      rounds_per_evaluation=FLAGS.rounds_per_eval,
      program_state_manager=program_state_manager,
      rounds_per_saving_program_state=FLAGS.rounds_per_checkpoint,
      metrics_managers=metrics_managers)

  test_metrics = federated_eval(state.model, [test_data])
  for metrics_manager in metrics_managers:
    metrics_manager.release(test_metrics, FLAGS.total_rounds + 1)


if __name__ == '__main__':
  app.run(main)

[((29, 5, 29, 37), 'utils.utils_impl.record_hparam_flags', 'utils_impl.record_hparam_flags', ({}, {}), '()', False, 'from utils import utils_impl\n'), ((31, 2, 31, 50), 'utils.optimizers.optimizer_utils.define_optimizer_flags', 'optimizer_utils.define_optimizer_flags', ({(31, 41, 31, 49): '"""client"""'}, {}), "('client')", False, 'from utils.optimizers import optimizer_utils\n'), ((32, 2, 32, 50), 'utils.optimizers.optimizer_utils.define_optimizer_flags', 'optimizer_utils.define_optimizer_flags', ({(32, 41, 32, 49): '"""server"""'}, {}), "('server')", False, 'from utils.optimizers import optimizer_utils\n'), ((34, 5, 34, 37), 'utils.utils_impl.record_hparam_flags', 'utils_impl.record_hparam_flags', ({}, {}), '()', False, 'from utils import utils_impl\n'), ((36, 2, 37, 75), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(36, 23, 36, 48): '"""client_epochs_per_round"""', (36, 50, 36, 51): '(1)', (37, 23, 37, 74): '"""Number of epochs in the client to take per round."""'}, {}), "('client_epochs_per_round', 1,\n    'Number of epochs in the client to take per round.')", False, 'from absl import flags\n'), ((38, 2, 38, 77), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(38, 23, 38, 42): '"""client_batch_size"""', (38, 44, 38, 46): '(20)', (38, 48, 38, 76): '"""Batch size on the clients."""'}, {}), "('client_batch_size', 20, 'Batch size on the clients.')", False, 'from absl import flags\n'), ((39, 2, 40, 63), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(39, 23, 39, 42): '"""clients_per_round"""', (39, 44, 39, 46): '(10)', (40, 23, 40, 62): '"""How many clients to sample per round."""'}, {}), "('clients_per_round', 10,\n    'How many clients to sample per round.')", False, 'from absl import flags\n'), ((41, 2, 42, 58), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(41, 23, 41, 52): '"""client_datasets_random_seed"""', (41, 54, 41, 55): '(1)', (42, 23, 42, 57): '"""Random seed for client sampling."""'}, {}), "('client_datasets_random_seed', 1,\n    'Random seed for client sampling.')", False, 'from absl import flags\n'), ((43, 2, 46, 37), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(44, 6, 44, 31): '"""max_elements_per_client"""', (44, 33, 44, 37): 'None', (44, 39, 46, 36): '"""Maximum number of elements for each training client. If set to None, all available examples are used."""'}, {}), "('max_elements_per_client', None,\n    'Maximum number of elements for each training client. If set to None, all available examples are used.'\n    )", False, 'from absl import flags\n'), ((49, 2, 49, 79), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(49, 23, 49, 37): '"""total_rounds"""', (49, 39, 49, 42): '(200)', (49, 44, 49, 78): '"""Number of total training rounds."""'}, {}), "('total_rounds', 200, 'Number of total training rounds.')", False, 'from absl import flags\n'), ((50, 2, 52, 58), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(51, 6, 51, 23): '"""experiment_name"""', (51, 25, 51, 29): 'None', (51, 31, 52, 57): '"""The name of this experiment. Will be append to --root_output_dir to separate experiment results."""'}, {}), "('experiment_name', None,\n    'The name of this experiment. Will be append to --root_output_dir to separate experiment results.'\n    )", False, 'from absl import flags\n'), ((53, 2, 54, 70), 'absl.flags.DEFINE_string', 'flags.DEFINE_string', ({(53, 22, 53, 39): '"""root_output_dir"""', (53, 41, 53, 56): '"""/tmp/fed_opt/"""', (54, 22, 54, 69): '"""Root directory for writing experiment output."""'}, {}), "('root_output_dir', '/tmp/fed_opt/',\n    'Root directory for writing experiment output.')", False, 'from absl import flags\n'), ((55, 2, 57, 74), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(56, 6, 56, 23): '"""rounds_per_eval"""', (56, 25, 56, 26): '(1)', (57, 6, 57, 73): '"""How often to evaluate the global model on the validation dataset."""'}, {}), "('rounds_per_eval', 1,\n    'How often to evaluate the global model on the validation dataset.')", False, 'from absl import flags\n'), ((58, 2, 61, 18), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(59, 6, 59, 31): '"""num_validation_examples"""', (59, 33, 59, 35): '(-1)', (59, 37, 61, 17): '"""The number of validationexamples to use. If set to -1, all available examples are used."""'}, {}), "('num_validation_examples', -1,\n    'The number of validationexamples to use. If set to -1, all available examples are used.'\n    )", False, 'from absl import flags\n'), ((62, 2, 63, 67), 'absl.flags.DEFINE_integer', 'flags.DEFINE_integer', ({(62, 23, 62, 46): '"""rounds_per_checkpoint"""', (62, 48, 62, 50): '(50)', (63, 23, 63, 66): '"""How often to checkpoint the global model."""'}, {}), "('rounds_per_checkpoint', 50,\n    'How often to checkpoint the global model.')", False, 'from absl import flags\n'), ((65, 5, 65, 37), 'utils.utils_impl.record_hparam_flags', 'utils_impl.record_hparam_flags', ({}, {}), '()', False, 'from utils import utils_impl\n'), ((67, 2, 69, 57), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(68, 6, 68, 12): '"""clip"""', (68, 14, 68, 18): 'None', (68, 20, 69, 56): '"""Clip value for fixed clipping or initial clip for adaptive clipping. If None, no clipping is used."""'}, {}), "('clip', None,\n    'Clip value for fixed clipping or initial clip for adaptive clipping. If None, no clipping is used.'\n    )", False, 'from absl import flags\n'), ((70, 2, 71, 77), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(70, 21, 70, 39): '"""noise_multiplier"""', (70, 41, 70, 45): 'None', (71, 21, 71, 76): '"""Noise multiplier. If None, non-DP aggregator is used."""'}, {}), "('noise_multiplier', None,\n    'Noise multiplier. If None, non-DP aggregator is used.')", False, 'from absl import flags\n'), ((72, 2, 74, 43), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(73, 6, 73, 35): '"""adaptive_clip_learning_rate"""', (73, 37, 73, 41): 'None', (73, 43, 74, 42): '"""Adaptive clip learning rate. If None, clip adaptation is not used."""'}, {}), "('adaptive_clip_learning_rate', None,\n    'Adaptive clip learning rate. If None, clip adaptation is not used.')", False, 'from absl import flags\n'), ((75, 2, 76, 50), 'absl.flags.DEFINE_float', 'flags.DEFINE_float', ({(75, 21, 75, 48): '"""target_unclipped_quantile"""', (75, 50, 75, 53): '(0.5)', (76, 21, 76, 49): '"""Target unclipped quantile."""'}, {}), "('target_unclipped_quantile', 0.5,\n    'Target unclipped quantile.')", False, 'from absl import flags\n'), ((77, 2, 78, 61), 'absl.flags.DEFINE_boolean', 'flags.DEFINE_boolean', ({(77, 23, 77, 42): '"""uniform_weighting"""', (77, 44, 77, 49): '(False)', (78, 23, 78, 60): '"""Whether to weigh clients uniformly."""'}, {}), "('uniform_weighting', False,\n    'Whether to weigh clients uniformly.')", False, 'from absl import flags\n'), ((81, 5, 81, 37), 'utils.utils_impl.record_hparam_flags', 'utils_impl.record_hparam_flags', ({}, {}), '()', False, 'from utils import utils_impl\n'), ((82, 2, 82, 32), 'utils.task_utils.define_task_flags', 'task_utils.define_task_flags', ({}, {}), '()', False, 'from utils import task_utils\n'), ((89, 16, 89, 59), 'utils.utils_impl.lookup_flag_values', 'utils_impl.lookup_flag_values', ({(89, 46, 89, 58): 'shared_flags'}, {}), '(shared_flags)', False, 'from utils import utils_impl\n'), ((92, 18, 92, 64), 'utils.utils_impl.lookup_flag_values', 'utils_impl.lookup_flag_values', ({(92, 48, 92, 63): 'optimizer_flags'}, {}), '(optimizer_flags)', False, 'from utils import utils_impl\n'), ((93, 18, 93, 78), 'utils.optimizers.optimizer_utils.remove_unused_flags', 'optimizer_utils.remove_unused_flags', ({(93, 54, 93, 62): '"""client"""', (93, 64, 93, 77): 'opt_flag_dict'}, {}), "('client', opt_flag_dict)", False, 'from utils.optimizers import optimizer_utils\n'), ((94, 18, 94, 78), 'utils.optimizers.optimizer_utils.remove_unused_flags', 'optimizer_utils.remove_unused_flags', ({(94, 54, 94, 62): '"""server"""', (94, 64, 94, 77): 'opt_flag_dict'}, {}), "('server', opt_flag_dict)", False, 'from utils.optimizers import optimizer_utils\n'), ((98, 19, 98, 60), 'utils.utils_impl.lookup_flag_values', 'utils_impl.lookup_flag_values', ({(98, 49, 98, 59): 'task_flags'}, {}), '(task_flags)', False, 'from utils import utils_impl\n'), ((100, 2, 101, 60), 'utils.training_utils.write_hparams_to_csv', 'training_utils.write_hparams_to_csv', ({(100, 38, 100, 49): 'hparam_dict', (100, 51, 100, 72): 'FLAGS.root_output_dir', (101, 38, 101, 59): 'FLAGS.experiment_name'}, {}), '(hparam_dict, FLAGS.root_output_dir,\n    FLAGS.experiment_name)', False, 'from utils import training_utils\n'), ((109, 24, 109, 80), 'utils.optimizers.optimizer_utils.create_optimizer_fn_from_flags', 'optimizer_utils.create_optimizer_fn_from_flags', ({(109, 71, 109, 79): '"""client"""'}, {}), "('client')", False, 'from utils.optimizers import optimizer_utils\n'), ((110, 24, 110, 80), 'utils.optimizers.optimizer_utils.create_optimizer_fn_from_flags', 'optimizer_utils.create_optimizer_fn_from_flags', ({(110, 71, 110, 79): '"""server"""'}, {}), "('server')", False, 'from utils.optimizers import optimizer_utils\n'), ((112, 22, 115, 49), 'tensorflow_federated.simulation.baselines.ClientSpec', 'tff.simulation.baselines.ClientSpec', (), '', True, 'import tensorflow_federated as tff\n'), ((116, 9, 116, 61), 'utils.task_utils.create_task_from_flags', 'task_utils.create_task_from_flags', ({(116, 43, 116, 60): 'train_client_spec'}, {}), '(train_client_spec)', False, 'from utils import task_utils\n'), ((118, 2, 118, 36), 'absl.logging.info', 'logging.info', ({(118, 15, 118, 35): '"""Trainable weights:"""'}, {}), "('Trainable weights:')", False, 'from absl import logging\n'), ((175, 22, 180, 59), 'tensorflow_federated.learning.build_federated_averaging_process', 'tff.learning.build_federated_averaging_process', (), '', True, 'import tensorflow_federated as tff\n'), ((184, 6, 185, 60), 'tensorflow_federated.simulation.compose_dataset_computation_with_iterative_process', 'tff.simulation.compose_dataset_computation_with_iterative_process', ({(185, 10, 185, 40): 'train_data.dataset_computation', (185, 42, 185, 59): 'iterative_process'}, {}), '(train_data\n    .dataset_computation, iterative_process)', True, 'import tensorflow_federated as tff\n'), ((194, 19, 194, 73), 'tensorflow_federated.learning.build_federated_evaluation', 'tff.learning.build_federated_evaluation', ({(194, 59, 194, 72): 'task.model_fn'}, {}), '(task.model_fn)', True, 'import tensorflow_federated as tff\n'), ((200, 44, 201, 51), 'utils.training_utils.create_managers', 'training_utils.create_managers', ({(201, 6, 201, 27): 'FLAGS.root_output_dir', (201, 29, 201, 50): 'FLAGS.experiment_name'}, {}), '(FLAGS.root_output_dir, FLAGS.experiment_name)', False, 'from utils import training_utils\n'), ((203, 10, 212, 40), 'tensorflow_federated.simulation.run_training_process', 'tff.simulation.run_training_process', (), '', True, 'import tensorflow_federated as tff\n'), ((220, 2, 220, 15), 'absl.app.run', 'app.run', ({(220, 10, 220, 14): 'main'}, {}), '(main)', False, 'from absl import app\n'), ((120, 4, 120, 66), 'absl.logging.info', 'logging.info', ({(120, 17, 120, 38): '"""name: %s  shape: %s"""', (120, 40, 120, 51): 'weight.name', (120, 53, 120, 65): 'weight.shape'}, {}), "('name: %s  shape: %s', weight.name, weight.shape)", False, 'from absl import logging\n'), ((188, 6, 189, 79), 'tensorflow_federated.simulation.build_uniform_sampling_fn', 'tff.simulation.build_uniform_sampling_fn', (), '', True, 'import tensorflow_federated as tff\n'), ((133, 28, 133, 67), 'tensorflow_federated.aggregators.UnweightedMeanFactory', 'tff.aggregators.UnweightedMeanFactory', ({}, {}), '()', True, 'import tensorflow_federated as tff\n'), ((135, 28, 135, 57), 'tensorflow_federated.aggregators.MeanFactory', 'tff.aggregators.MeanFactory', ({}, {}), '()', True, 'import tensorflow_federated as tff\n'), ((149, 28, 150, 36), 'tensorflow_federated.aggregators.clipping_factory', 'tff.aggregators.clipping_factory', ({(150, 10, 150, 14): 'clip', (150, 16, 150, 35): 'aggregation_factory'}, {}), '(clip, aggregation_factory)', True, 'import tensorflow_federated as tff\n'), ((160, 28, 163, 26), 'tensorflow_federated.aggregators.DifferentiallyPrivateFactory.gaussian_fixed', 'tff.aggregators.DifferentiallyPrivateFactory.gaussian_fixed', (), '', True, 'import tensorflow_federated as tff\n'), ((168, 28, 173, 58), 'tensorflow_federated.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive', 'tff.aggregators.DifferentiallyPrivateFactory.gaussian_adaptive', (), '', True, 'import tensorflow_federated as tff\n'), ((145, 15, 148, 60), 'tensorflow_federated.aggregators.PrivateQuantileEstimationProcess.no_noise', 'tff.aggregators.PrivateQuantileEstimationProcess.no_noise', (), '', True, 'import tensorflow_federated as tff\n'), ((127, 21, 127, 60), 'tensorflow.squeeze', 'tf.squeeze', ({(127, 32, 127, 59): "local_outputs['num_tokens']"}, {}), "(local_outputs['num_tokens'])", True, 'import tensorflow as tf\n')]