Python absl.flags.DEFINE_string() Examples

def define_flags(specs=None):
  """Define a command line flag for each ParamSpec in flags.param_specs."""
  specs = specs or param_specs
  define_flag = {
      'boolean': absl_flags.DEFINE_boolean,
      'float': absl_flags.DEFINE_float,
      'integer': absl_flags.DEFINE_integer,
      'string': absl_flags.DEFINE_string,
      'enum': absl_flags.DEFINE_enum,
      'list': absl_flags.DEFINE_list
  }
  for name, param_spec in six.iteritems(specs):
    if param_spec.flag_type not in define_flag:
      raise ValueError('Unknown flag_type %s' % param_spec.flag_type)
    else:
      define_flag[param_spec.flag_type](name, param_spec.default_value,
                                        help=param_spec.description,
                                        **param_spec.kwargs) 

def define_compute_bleu_flags():
  """Add flags for computing BLEU score."""
  flags.DEFINE_string(
      name="translation", default=None,
      help=flags_core.help_wrap("File containing translated text."))
  flags.mark_flag_as_required("translation")

  flags.DEFINE_string(
      name="reference", default=None,
      help=flags_core.help_wrap("File containing reference translation."))
  flags.mark_flag_as_required("reference")

  flags.DEFINE_enum(
      name="bleu_variant", short_name="bv", default="both",
      enum_values=["both", "uncased", "cased"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Specify one or more BLEU variants to calculate. Variants: \"cased\""
          ", \"uncased\", or \"both\".")) 

def _define_flags(options):
  for option in _filter_options(options):
    name, type_, default_ = option[:3]
    if type_ == 'boolean':
      flags.DEFINE_boolean(name, default_, '')
    elif type_ == 'integer':
      flags.DEFINE_integer(name, default_, '')
    elif type_ == 'float':
      flags.DEFINE_float(name, default_, '')
    elif type_ == 'string':
      flags.DEFINE_string(name, default_, '')
    elif type_ == 'list_of_ints':
      flags.DEFINE_string(name, default_, '')
    else:
      assert 'Unexpected option type %s' % type_


# Define command-line flags for all options (unless `internal`). 

def initialize_common_flags():
  """Define the common flags across models."""
  define_common_hparams_flags()

  flags_core.define_device(tpu=True)
  flags_core.define_base(
      num_gpu=True, model_dir=False, data_dir=False, batch_size=False)
  flags_core.define_distribution(worker_hosts=True, task_index=True)
  flags_core.define_performance(all_reduce_alg=True, num_packs=True)

  # Reset the default value of num_gpus to zero.
  FLAGS.num_gpus = 0

  flags.DEFINE_string(
      'strategy_type', 'mirrored', 'Type of distribute strategy.'
      'One of mirrored, tpu and multiworker.') 

def define_flags():
  """Define a command line flag for each ParamSpec in flags.param_specs."""
  define_flag = {
      'boolean': absl_flags.DEFINE_boolean,
      'float': absl_flags.DEFINE_float,
      'integer': absl_flags.DEFINE_integer,
      'string': absl_flags.DEFINE_string,
      'enum': absl_flags.DEFINE_enum,
      'list': absl_flags.DEFINE_list
  }
  for name, param_spec in six.iteritems(param_specs):
    if param_spec.flag_type not in define_flag:
      raise ValueError('Unknown flag_type %s' % param_spec.flag_type)
    else:
      define_flag[param_spec.flag_type](name, param_spec.default_value,
                                        help=param_spec.description,
                                        **param_spec.kwargs) 

def define_compute_bleu_flags():
  """Add flags for computing BLEU score."""
  flags.DEFINE_string(
      name="translation", default=None,
      help=flags_core.help_wrap("File containing translated text."))
  flags.mark_flag_as_required("translation")

  flags.DEFINE_string(
      name="reference", default=None,
      help=flags_core.help_wrap("File containing reference translation."))
  flags.mark_flag_as_required("reference")

  flags.DEFINE_enum(
      name="bleu_variant", short_name="bv", default="both",
      enum_values=["both", "uncased", "cased"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Specify one or more BLEU variants to calculate. Variants: \"cased\""
          ", \"uncased\", or \"both\".")) 

def define_compute_bleu_flags():
  """Add flags for computing BLEU score."""
  flags.DEFINE_string(
      name="translation", default=None,
      help=flags_core.help_wrap("File containing translated text."))
  flags.mark_flag_as_required("translation")

  flags.DEFINE_string(
      name="reference", default=None,
      help=flags_core.help_wrap("File containing reference translation."))
  flags.mark_flag_as_required("reference")

  flags.DEFINE_enum(
      name="bleu_variant", short_name="bv", default="both",
      enum_values=["both", "uncased", "cased"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Specify one or more BLEU variants to calculate. Variants: \"cased\""
          ", \"uncased\", or \"both\".")) 

def config_with_absl(self):
    # Run this before calling `app.run(main)` etc
    import absl.flags as absl_FLAGS
    from absl import app, flags as absl_flags

    self.use_absl = True
    self.absl_flags = absl_flags
    absl_defs = { bool: absl_flags.DEFINE_bool,
                  int:  absl_flags.DEFINE_integer,
                  str:  absl_flags.DEFINE_string,
                  'enum': absl_flags.DEFINE_enum }

    for name, val in self.values.items():
      flag_type, meta_args, meta_kwargs = self.meta[name]
      absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)

    app.call_after_init(lambda: self.complete_absl_config(absl_flags)) 

def record_new_flags() -> Iterator[List[str]]:
  """A context manager that returns all flags created in it's scope.

  This is useful to define all of the flags which should be considered
  hyperparameters of the training run, without needing to repeat them.

  Example usage:
  ```python
  with record_new_flags() as hparam_flags:
      flags.DEFINE_string('exp_name', 'name', 'Unique name for the experiment.')
  ```

  Check `research/emnist/run_experiment.py` for more details about the usage.

  Yields:
    A list of all newly created flags.
  """
  old_flags = set(iter(flags.FLAGS))
  new_flags = []
  yield new_flags
  new_flags.extend([f for f in flags.FLAGS if f not in old_flags]) 

def define_compute_bleu_flags():
  """Add flags for computing BLEU score."""
  flags.DEFINE_string(
      name="translation", default=None,
      help=flags_core.help_wrap("File containing translated text."))
  flags.mark_flag_as_required("translation")

  flags.DEFINE_string(
      name="reference", default=None,
      help=flags_core.help_wrap("File containing reference translation."))
  flags.mark_flag_as_required("reference")

  flags.DEFINE_enum(
      name="bleu_variant", short_name="bv", default="both",
      enum_values=["both", "uncased", "cased"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Specify one or more BLEU variants to calculate. Variants: \"cased\""
          ", \"uncased\", or \"both\".")) 

def define_compute_bleu_flags():
  """Add flags for computing BLEU score."""
  flags.DEFINE_string(
      name="translation", default=None,
      help=flags_core.help_wrap("File containing translated text."))
  flags.mark_flag_as_required("translation")

  flags.DEFINE_string(
      name="reference", default=None,
      help=flags_core.help_wrap("File containing reference translation."))
  flags.mark_flag_as_required("reference")

  flags.DEFINE_enum(
      name="bleu_variant", short_name="bv", default="both",
      enum_values=["both", "uncased", "cased"], case_sensitive=False,
      help=flags_core.help_wrap(
          "Specify one or more BLEU variants to calculate. Variants: \"cased\""
          ", \"uncased\", or \"both\".")) 

def test_helpfull_message(self):
    flags.DEFINE_string(
        'non_main_module_flag', 'default', 'help',
        module_name='other.module', flag_values=self._absl_flags)
    parser = argparse_flags.ArgumentParser(
        inherited_absl_flags=self._absl_flags)
    with self.assertRaises(SystemExit),\
        mock.patch.object(sys, 'stdout', new=six.StringIO()) as mock_stdout:
      parser.parse_args(['--helpfull'])
    stdout_message = mock_stdout.getvalue()
    logging.info('captured stdout message:\n%s', stdout_message)
    self.assertIn('--non_main_module_flag', stdout_message)
    self.assertIn('other.module', stdout_message)
    # Make sure the main module is not included.
    self.assertNotIn(sys.argv[0], stdout_message)
    # Special flags defined in absl.flags.
    self.assertIn('absl.flags:', stdout_message)
    self.assertIn('--flagfile', stdout_message)
    self.assertIn('--undefok', stdout_message) 

def define_flags(flag_values=FLAGS):
  """Defines some flags.

  Args:
    flag_values: The FlagValues object we want to register the flags
      with.
  """
  # The 'tmod_bar_' prefix (short for 'test_module_bar') ensures there
  # is no name clash with the existing flags.
  flags.DEFINE_boolean('tmod_bar_x', True, 'Boolean flag.',
                       flag_values=flag_values)
  flags.DEFINE_string('tmod_bar_y', 'default', 'String flag.',
                      flag_values=flag_values)
  flags.DEFINE_boolean('tmod_bar_z', False,
                       'Another boolean flag from module bar.',
                       flag_values=flag_values)
  flags.DEFINE_integer('tmod_bar_t', 4, 'Sample int flag.',
                       flag_values=flag_values)
  flags.DEFINE_integer('tmod_bar_u', 5, 'Sample int flag.',
                       flag_values=flag_values)
  flags.DEFINE_integer('tmod_bar_v', 6, 'Sample int flag.',
                       flag_values=flag_values) 

def define_common_tpu_flags():
  """Define the flags related to TPU's."""
  flags.DEFINE_string(
      'tpu', default=None,
      help='The Cloud TPU to use for training. This should be either the name '
      'used when creating the Cloud TPU, or a grpc://ip.address.of.tpu:8470 '
      'url.')

  flags.DEFINE_string(
      'gcp_project', default=None,
      help='Project name for the Cloud TPU-enabled project. If not specified, we '
      'will attempt to automatically detect the GCE project from metadata.')

  flags.DEFINE_string(
      'tpu_zone', default=None,
      help='GCE zone where the Cloud TPU is located in. If not specified, we '
      'will attempt to automatically detect the GCE project from metadata.')

  flags.DEFINE_string(
      'eval_master', default='',
      help='GRPC URL of the eval master. Set to an appropiate value when running '
      'on CPU/GPU.') 

def setUp(self):
    self._absl_flags = flags.FlagValues()
    flags.DEFINE_bool(
        'absl_bool', None, 'help for --absl_bool.',
        short_name='b', flag_values=self._absl_flags)
    # Add a boolean flag that starts with "no", to verify it can correctly
    # handle the "no" prefixes in boolean flags.
    flags.DEFINE_bool(
        'notice', None, 'help for --notice.',
        flag_values=self._absl_flags)
    flags.DEFINE_string(
        'absl_string', 'default', 'help for --absl_string=%.',
        short_name='s', flag_values=self._absl_flags)
    flags.DEFINE_integer(
        'absl_integer', 1, 'help for --absl_integer.',
        flag_values=self._absl_flags)
    flags.DEFINE_float(
        'absl_float', 1, 'help for --absl_integer.',
        flag_values=self._absl_flags)
    flags.DEFINE_enum(
        'absl_enum', 'apple', ['apple', 'orange'], 'help for --absl_enum.',
        flag_values=self._absl_flags) 

def test_help_non_main_module_key_flags(self):
    flags.DEFINE_string(
        'non_main_module_flag', 'default', 'help',
        module_name='other.module', flag_values=self._absl_flags)
    flags.declare_key_flag('non_main_module_flag', flag_values=self._absl_flags)
    parser = argparse_flags.ArgumentParser(
        inherited_absl_flags=self._absl_flags)
    help_message = parser.format_help()

    # Main module key fags are printed in the help message, even if the flag
    # is defined in another module.
    self.assertIn('non_main_module_flag', help_message) 

def define_mnist_eager_flags():
  """Defined flags and defaults for MNIST in eager mode."""
  flags_core.define_base_eager(clean=True, train_epochs=True)
  flags_core.define_image()
  flags.adopt_module_key_flags(flags_core)

  flags.DEFINE_integer(
      name='log_interval', short_name='li', default=10,
      help=flags_core.help_wrap('batches between logging training status'))

  flags.DEFINE_string(
      name='output_dir', short_name='od', default=None,
      help=flags_core.help_wrap('Directory to write TensorBoard summaries'))

  flags.DEFINE_float(name='learning_rate', short_name='lr', default=0.01,
                     help=flags_core.help_wrap('Learning rate.'))

  flags.DEFINE_float(name='momentum', short_name='m', default=0.5,
                     help=flags_core.help_wrap('SGD momentum.'))

  flags.DEFINE_bool(name='no_gpu', short_name='nogpu', default=False,
                    help=flags_core.help_wrap(
                        'disables GPU usage even if a GPU is available'))

  flags_core.set_defaults(
      data_dir='/tmp/tensorflow/mnist/input_data',
      model_dir='/tmp/tensorflow/mnist/checkpoints/',
      batch_size=100,
      train_epochs=10,
  ) 

def define_distribution(worker_hosts=True, task_index=True):
  """Register distributed execution flags.

  Args:
    worker_hosts: Create a flag for specifying comma-separated list of workers.
    task_index: Create a flag for specifying index of task.

  Returns:
    A list of flags for core.py to marks as key flags.
  """
  key_flags = []

  if worker_hosts:
    flags.DEFINE_string(
        name='worker_hosts', default=None,
        help=help_wrap(
            'Comma-separated list of worker ip:port pairs for running '
            'multi-worker models with DistributionStrategy.  The user would '
            'start the program on each host with identical value for this '
            'flag.'))

  if task_index:
    flags.DEFINE_integer(
        name='task_index', default=-1,
        help=help_wrap('If multi-worker training, the task_index of this '
                       'worker.'))

  return key_flags 

def define_mnist_eager_flags():
  """Defined flags and defaults for MNIST in eager mode."""
  flags_core.define_base_eager()
  flags_core.define_image()
  flags.adopt_module_key_flags(flags_core)

  flags.DEFINE_integer(
      name='log_interval', short_name='li', default=10,
      help=flags_core.help_wrap('batches between logging training status'))

  flags.DEFINE_string(
      name='output_dir', short_name='od', default=None,
      help=flags_core.help_wrap('Directory to write TensorBoard summaries'))

  flags.DEFINE_float(name='learning_rate', short_name='lr', default=0.01,
                     help=flags_core.help_wrap('Learning rate.'))

  flags.DEFINE_float(name='momentum', short_name='m', default=0.5,
                     help=flags_core.help_wrap('SGD momentum.'))

  flags.DEFINE_bool(name='no_gpu', short_name='nogpu', default=False,
                    help=flags_core.help_wrap(
                        'disables GPU usage even if a GPU is available'))

  flags_core.set_defaults(
      data_dir='/tmp/tensorflow/mnist/input_data',
      model_dir='/tmp/tensorflow/mnist/checkpoints/',
      batch_size=100,
      train_epochs=10,
  ) 

def DEFINE_string(name, default, help):  # pylint: disable=invalid-name,redefined-builtin
  param_specs[name] = ParamSpec('string', default, help, {}) 

def define_data_download_flags():
  """Add flags specifying data download arguments."""
  flags.DEFINE_string(
      name="data_dir", short_name="dd", default="/tmp/translate_ende",
      help=flags_core.help_wrap(
          "Directory for where the translate_ende_wmt32k dataset is saved."))
  flags.DEFINE_string(
      name="raw_dir", short_name="rd", default="/tmp/translate_ende_raw",
      help=flags_core.help_wrap(
          "Path where the raw data will be downloaded and extracted."))
  flags.DEFINE_bool(
      name="search", default=False,
      help=flags_core.help_wrap(
          "If set, use binary search to find the vocabulary set with size"
          "closest to the target size (%d)." % _TARGET_VOCAB_SIZE)) 

def define_translate_flags():
  """Define flags used for translation script."""
  # Model flags
  flags.DEFINE_string(
      name="model_dir", short_name="md", default="/tmp/transformer_model",
      help=flags_core.help_wrap(
          "Directory containing Transformer model checkpoints."))
  flags.DEFINE_enum(
      name="param_set", short_name="mp", default="big",
      enum_values=["base", "big"],
      help=flags_core.help_wrap(
          "Parameter set to use when creating and training the model. The "
          "parameters define the input shape (batch size and max length), "
          "model configuration (size of embedding, # of hidden layers, etc.), "
          "and various other settings. The big parameter set increases the "
          "default batch size, embedding/hidden size, and filter size. For a "
          "complete list of parameters, please see model/model_params.py."))
  flags.DEFINE_string(
      name="vocab_file", short_name="vf", default=None,
      help=flags_core.help_wrap(
          "Path to subtoken vocabulary file. If data_download.py was used to "
          "download and encode the training data, look in the data_dir to find "
          "the vocab file."))
  flags.mark_flag_as_required("vocab_file")

  flags.DEFINE_string(
      name="text", default=None,
      help=flags_core.help_wrap(
          "Text to translate. Output will be printed to console."))
  flags.DEFINE_string(
      name="file", default=None,
      help=flags_core.help_wrap(
          "File containing text to translate. Translation will be printed to "
          "console and, if --file_out is provided, saved to an output file."))
  flags.DEFINE_string(
      name="file_out", default=None,
      help=flags_core.help_wrap(
          "If --file flag is specified, save translation to this file.")) 

def define_mnist_eager_flags():
  """Defined flags and defaults for MNIST in eager mode."""
  flags_core.define_base_eager()
  flags_core.define_image()
  flags.adopt_module_key_flags(flags_core)

  flags.DEFINE_integer(
      name='log_interval', short_name='li', default=10,
      help=flags_core.help_wrap('batches between logging training status'))

  flags.DEFINE_string(
      name='output_dir', short_name='od', default=None,
      help=flags_core.help_wrap('Directory to write TensorBoard summaries'))

  flags.DEFINE_float(name='learning_rate', short_name='lr', default=0.01,
                     help=flags_core.help_wrap('Learning rate.'))

  flags.DEFINE_float(name='momentum', short_name='m', default=0.5,
                     help=flags_core.help_wrap('SGD momentum.'))

  flags.DEFINE_bool(name='no_gpu', short_name='nogpu', default=False,
                    help=flags_core.help_wrap(
                        'disables GPU usage even if a GPU is available'))

  flags_core.set_defaults(
      data_dir='/tmp/tensorflow/mnist/input_data',
      model_dir='/tmp/tensorflow/mnist/checkpoints/',
      batch_size=100,
      train_epochs=10,
  ) 

def define_device(tpu=True):
  """Register device specific flags.
  Args:
    tpu: Create flags to specify TPU operation.
  Returns:
    A list of flags for core.py to marks as key flags.
  """

  key_flags = []

  if tpu:
    flags.DEFINE_string(
        name="tpu", default=None,
        help=help_wrap(
            "The Cloud TPU to use for training. This should be either the name "
            "used when creating the Cloud TPU, or a "
            "grpc://ip.address.of.tpu:8470 url. Passing `local` will use the"
            "CPU of the local instance instead. (Good for debugging.)"))
    key_flags.append("tpu")

    flags.DEFINE_string(
        name="tpu_zone", default=None,
        help=help_wrap(
            "[Optional] GCE zone where the Cloud TPU is located in. If not "
            "specified, we will attempt to automatically detect the GCE "
            "project from metadata."))

    flags.DEFINE_string(
        name="tpu_gcp_project", default=None,
        help=help_wrap(
            "[Optional] Project name for the Cloud TPU-enabled project. If not "
            "specified, we will attempt to automatically detect the GCE "
            "project from metadata."))

    flags.DEFINE_integer(name="num_tpu_shards", default=8,
                         help=help_wrap("Number of shards (TPU chips)."))

  return key_flags 

def test_help_non_main_module_flags(self):
    flags.DEFINE_string(
        'non_main_module_flag', 'default', 'help',
        module_name='other.module', flag_values=self._absl_flags)
    parser = argparse_flags.ArgumentParser(
        inherited_absl_flags=self._absl_flags)
    help_message = parser.format_help()

    # Non main module key flags are not printed in the help message.
    self.assertNotIn('non_main_module_flag', help_message) 

def DEFINE_string(name, default, help):  # pylint: disable=invalid-name,redefined-builtin
  param_specs[name] = ParamSpec('string', default, help, {}) 

def test_flag_help_in_xml_string(self):
    flags.DEFINE_string('file_path', '/path/to/my/dir', 'A test string flag.',
                        flag_values=self.fv)
    expected_output = (
        '<flag>\n'
        '  <file>simple_module</file>\n'
        '  <name>file_path</name>\n'
        '  <meaning>A test string flag.</meaning>\n'
        '  <default>/path/to/my/dir</default>\n'
        '  <current>/path/to/my/dir</current>\n'
        '  <type>string</type>\n'
        '</flag>\n')
    self._check_flag_help_in_xml('file_path', 'simple_module', expected_output) 

def test_nonglobal_flags(self):
    """Test use of non-global FlagValues."""
    nonglobal_flags = flags.FlagValues()
    flags.DEFINE_string('nonglobal_flag', 'Bob', 'flaghelp', nonglobal_flags)
    argv = ('./program',
            '--nonglobal_flag=Mary',
            'extra')
    argv = nonglobal_flags(argv)
    self.assertEqual(len(argv), 2, 'wrong number of arguments pulled')
    self.assertEqual(argv[0], './program', 'program name not preserved')
    self.assertEqual(argv[1], 'extra', 'extra argument not preserved')
    self.assertEqual(nonglobal_flags['nonglobal_flag'].value, 'Mary') 

def define_data_download_flags():
  """Add flags specifying data download arguments."""
  flags.DEFINE_string(
      name="data_dir", short_name="dd", default="/tmp/translate_ende",
      help=flags_core.help_wrap(
          "Directory for where the translate_ende_wmt32k dataset is saved."))
  flags.DEFINE_string(
      name="raw_dir", short_name="rd", default="/tmp/translate_ende_raw",
      help=flags_core.help_wrap(
          "Path where the raw data will be downloaded and extracted."))
  flags.DEFINE_bool(
      name="search", default=False,
      help=flags_core.help_wrap(
          "If set, use binary search to find the vocabulary set with size"
          "closest to the target size (%d)." % _TARGET_VOCAB_SIZE)) 

def test_one_dash_arg_first(self):
    flags.DEFINE_string('onedash_name', 'Bob', 'namehelp',
                        flag_values=self.flag_values)
    flags.DEFINE_string('onedash_blame', 'Rob', 'blamehelp',
                        flag_values=self.flag_values)
    argv = ('./program',
            '-',
            '--onedash_name=Harry')
    with _use_gnu_getopt(self.flag_values, False):
      argv = self.flag_values(argv)
      self.assertEqual(len(argv), 3)
      self.assertEqual(argv[1], '-')
      self.assertEqual(argv[2], '--onedash_name=Harry')