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

"""A registry of experiment configurations."""
import collections
import functools
import re
from typing import List, Mapping, Optional

from paxml import base_experiment

BaseExperimentT = base_experiment.BaseExperimentT


class _ExperimentRegistryHelper:
  """Helper class encapsulating a global registry keyed by experiment name."""

  # Global variable for the experiment configuration registry
  # A mapping from a canonical key to the BaseExperimentT class.
  _registry = {}
  # A mapping from a secondary key to all matching canonical keys.
  _secondary_keys = collections.defaultdict(list)

  @classmethod
  def custom_secondary_keys(cls, canonical_key) -> List[str]:
    """Returns the list of custom secondary keys."""
    ret = []
    parts = canonical_key.split('.')
    try:
      idx = parts.index('params')
      if idx > 0:
        del parts[idx]
        new_key = '.'.join(parts[idx - 1:])
        ret.append(new_key)
    except ValueError:
      pass
    return ret

  @classmethod
  def register(cls,
               experiment_class: Optional[BaseExperimentT] = None,
               *,
               allow_overwrite=False):
    """Registers an experiment configuration class into the global registry.

    If allow_overwrite is True, repeated registering of an existing class
    refreshes its definition in the global registry. Otherwise re-registering
    raises an error.

    Usage example:
      @experiment_registry.register
      class MyExperiment(base_experiment.BaseExperiment):
        pass

      # Allows re-registering.
      @experiment_registry.register(allow_overwrite=True)
      class MyExperiment(base_experiment.BaseExperiment):
        pass

    Args:
      experiment_class: a BaseExperimentT class.
      allow_overwrite: bool, whether re-register an existing class is allowed.

    Returns:
      experiment_class itself, so this can be used as a class decorator, or
      a decorator.

    Raises:
      ValueError: if allow_overwrite is False and the same class (same module
      and class name) has already been registered.
    """
    if experiment_class is None:
      # decorated with @register(...)
      return functools.partial(cls.register, allow_overwrite=allow_overwrite)

    # canonical key is the full path.
    canonical_key = (
        experiment_class.__module__ + '.' + experiment_class.__name__)
    preexisting = canonical_key in cls._registry
    if preexisting and not allow_overwrite:
      raise ValueError(f'Experiment already registered: {canonical_key}')
    cls._registry[canonical_key] = experiment_class
    print('Registered experiment `%s`%s' % (
        canonical_key, ' (overwritten)' if preexisting else ''))
    if preexisting:
      # No need to update secondary keys.
      return experiment_class
    # Add secondary keys, which can be any partial paths.
    secondary_keys = set()
    parts = canonical_key.split('.')
    for i in range(len(parts)):
      # Any partial path is a valid secondary key.
      secondary_keys.add('.'.join(parts[i:]))

    for k in cls.custom_secondary_keys(canonical_key):
      secondary_keys.add(k)
    for k in secondary_keys:
      cls._secondary_keys[k].append(canonical_key)
    return experiment_class

  @classmethod
  def get(cls, key: str) -> Optional[BaseExperimentT]:
    """Retrieves an experiment from the global registry from the input key.

    Args:
      key: string, a (secondary) key for the experiment. Any partial path works
        as long as it uniquely identifies the experiment.

    Returns:
      None if no match is found, or the unique BaseExperimentT that
      matches the provided key.

    Raises:
      If the provided key does not uniquely identifies an experiment.
    """
    canonical_keys = cls._secondary_keys.get(key)
    if not canonical_keys:
      return None
    if len(canonical_keys) > 1:
      matches = ', '.join(canonical_keys)
      raise ValueError(f'key={key} does not uniquely identify an experiment. '
                       f'possible matches: {matches}')
    return cls._registry.get(canonical_keys[0])

  @classmethod
  def get_docstring_tags(cls, key: str) -> Optional[Mapping[str, str]]:
    """Retrieves @tags from the class docstring."""
    params = cls.get(key)
    if not params:
      # Note that None means no experiment found, different from an empty map.
      return None
    ret = {}
    for line in params.__doc__.split('\n'):
      m = re.search(r'@(.*): (.*)', line)
      if m:
        tag, value = m.group(1), m.group(2)
        ret[tag] = value
    return ret

  @classmethod
  def get_all(cls) -> Mapping[str, BaseExperimentT]:
    """Retrieves all the experiment configurations from the global registry.

    Returns:
      A dict of {experiment canonical key, experiment class}.
    """
    return cls._registry


register = _ExperimentRegistryHelper.register

# Unless you are writing a binary target like main.py that is part of
# the BUILD macro pax_targets(), you should not be calling get().
# Explicitly import the module that defines the experiment config like a regular
# Python class instead.
get = _ExperimentRegistryHelper.get

get_all = _ExperimentRegistryHelper.get_all
get_docstring_tags = _ExperimentRegistryHelper.get_docstring_tags
<commit_msg>Print the registration message to stderr instead of stdout.<commit_after># coding=utf-8
# Copyright 2022 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.

"""A registry of experiment configurations."""
import collections
import functools
import re
import sys
from typing import List, Mapping, Optional

from paxml import base_experiment

BaseExperimentT = base_experiment.BaseExperimentT


class _ExperimentRegistryHelper:
  """Helper class encapsulating a global registry keyed by experiment name."""

  # Global variable for the experiment configuration registry
  # A mapping from a canonical key to the BaseExperimentT class.
  _registry = {}
  # A mapping from a secondary key to all matching canonical keys.
  _secondary_keys = collections.defaultdict(list)

  @classmethod
  def custom_secondary_keys(cls, canonical_key) -> List[str]:
    """Returns the list of custom secondary keys."""
    ret = []
    parts = canonical_key.split('.')
    try:
      idx = parts.index('params')
      if idx > 0:
        del parts[idx]
        new_key = '.'.join(parts[idx - 1:])
        ret.append(new_key)
    except ValueError:
      pass
    return ret

  @classmethod
  def register(cls,
               experiment_class: Optional[BaseExperimentT] = None,
               *,
               allow_overwrite=False):
    """Registers an experiment configuration class into the global registry.

    If allow_overwrite is True, repeated registering of an existing class
    refreshes its definition in the global registry. Otherwise re-registering
    raises an error.

    Usage example:
      @experiment_registry.register
      class MyExperiment(base_experiment.BaseExperiment):
        pass

      # Allows re-registering.
      @experiment_registry.register(allow_overwrite=True)
      class MyExperiment(base_experiment.BaseExperiment):
        pass

    Args:
      experiment_class: a BaseExperimentT class.
      allow_overwrite: bool, whether re-register an existing class is allowed.

    Returns:
      experiment_class itself, so this can be used as a class decorator, or
      a decorator.

    Raises:
      ValueError: if allow_overwrite is False and the same class (same module
      and class name) has already been registered.
    """
    if experiment_class is None:
      # decorated with @register(...)
      return functools.partial(cls.register, allow_overwrite=allow_overwrite)

    # canonical key is the full path.
    canonical_key = (
        experiment_class.__module__ + '.' + experiment_class.__name__)
    preexisting = canonical_key in cls._registry
    if preexisting and not allow_overwrite:
      raise ValueError(f'Experiment already registered: {canonical_key}')
    cls._registry[canonical_key] = experiment_class
    print(
        'Registered experiment `%s`%s' %
        (canonical_key, ' (overwritten)' if preexisting else ''),
        file=sys.stderr)
    if preexisting:
      # No need to update secondary keys.
      return experiment_class
    # Add secondary keys, which can be any partial paths.
    secondary_keys = set()
    parts = canonical_key.split('.')
    for i in range(len(parts)):
      # Any partial path is a valid secondary key.
      secondary_keys.add('.'.join(parts[i:]))

    for k in cls.custom_secondary_keys(canonical_key):
      secondary_keys.add(k)
    for k in secondary_keys:
      cls._secondary_keys[k].append(canonical_key)
    return experiment_class

  @classmethod
  def get(cls, key: str) -> Optional[BaseExperimentT]:
    """Retrieves an experiment from the global registry from the input key.

    Args:
      key: string, a (secondary) key for the experiment. Any partial path works
        as long as it uniquely identifies the experiment.

    Returns:
      None if no match is found, or the unique BaseExperimentT that
      matches the provided key.

    Raises:
      If the provided key does not uniquely identifies an experiment.
    """
    canonical_keys = cls._secondary_keys.get(key)
    if not canonical_keys:
      return None
    if len(canonical_keys) > 1:
      matches = ', '.join(canonical_keys)
      raise ValueError(f'key={key} does not uniquely identify an experiment. '
                       f'possible matches: {matches}')
    return cls._registry.get(canonical_keys[0])

  @classmethod
  def get_docstring_tags(cls, key: str) -> Optional[Mapping[str, str]]:
    """Retrieves @tags from the class docstring."""
    params = cls.get(key)
    if not params:
      # Note that None means no experiment found, different from an empty map.
      return None
    ret = {}
    for line in params.__doc__.split('\n'):
      m = re.search(r'@(.*): (.*)', line)
      if m:
        tag, value = m.group(1), m.group(2)
        ret[tag] = value
    return ret

  @classmethod
  def get_all(cls) -> Mapping[str, BaseExperimentT]:
    """Retrieves all the experiment configurations from the global registry.

    Returns:
      A dict of {experiment canonical key, experiment class}.
    """
    return cls._registry


register = _ExperimentRegistryHelper.register

# Unless you are writing a binary target like main.py that is part of
# the BUILD macro pax_targets(), you should not be calling get().
# Explicitly import the module that defines the experiment config like a regular
# Python class instead.
get = _ExperimentRegistryHelper.get

get_all = _ExperimentRegistryHelper.get_all
get_docstring_tags = _ExperimentRegistryHelper.get_docstring_tags

{'fraction_non_alphanumeric': 0.07986463620981388, 'fraction_numerical': 0.0025380710659898475, 'mean_word_length': 3.2460488505747125, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 2}, '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': '6173620', 'n_tokens_mistral': 3262, 'n_tokens_neox': 3196, 'n_words': 1176}