<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.

"""Base classes for the Praxis input layers."""

from __future__ import annotations

import abc
import copy
import re
from typing import Any, Callable, Dict, Iterable, Optional, Sequence

from absl import logging
import jax
from lingvo.core import cluster_factory
from lingvo.core import datasource
import numpy as np
from praxis import base_hyperparams
from praxis import py_utils
from praxis import pytypes
import tensorflow.compat.v2 as tf

NestedMap = py_utils.NestedMap
NestedJTensor = pytypes.NestedJTensor
Nested = pytypes.Nested
NestedShapeDtypeStruct = pytypes.NestedShapeDtypeStruct
instantiate = base_hyperparams.instantiate


class BaseInput(base_hyperparams.BaseParameterizable):
  """Base class for Jax input classes.

  During paxml's train, on each host an input instance will be
  created (instantiate(input_p)), and then get_next() is iteratively
  called in eager mode to generate one batch of data for each step
  of train/eval/etc.

  If supported, for eval, reset() is called after each eval step.
  See p.reset_for_eval below.

  A tf.data based input should inherit this class directly and implement
  get_next() and reset(). For an example of how to handle sharding for both
  training and eval data, please refer to the implementation of
  TFRecordBertInput at tasks/lm/input_generator.py.

  If there is already an Lingvo TF input generator that one would like to
  use directly, please use LingvoInputAdaptor below.
  """
  _VALIDATE_BATCH_SIZE_NOT_NONE = True

  class HParams(base_hyperparams.BaseParameterizable.HParams):
    """Hyper-parameters associated with this Input class.

    Attributes:
      batch_size: The (Jax per process) Batch size. Each call to get_next()
        returns a batch with this batch size.
      num_infeed_hosts: Usually set to jax.process_count(). Implementation must
        ensure that the data is sharded into this many shard.
      infeed_host_index: Usually set to jax.process_index(). Implementation must
        ensure that each instance returns a shard with this index.
      input_random_seed: If set, implementation must ensure that this is used to
        seed randomness, e.g. when shuffling in a deterministic manner.
      reset_for_eval: If set, eval will continue until tf.errors.OutOfRange is
        raised, and reset() will called for each eval. Implementation must
        ensure that all variant p.infeed_host_index instances raise after the
        same numberof calls to get_next() to ensure synchronization across
        hosts. If notset, get_next() must never raise.
      eval_loop_num_batches: Num of batches to process per eval loop. Must be >=
        1. This valueis ignored if reset_for_eval is set True, in which case,
        this value is dynamically determined by the number of available batches.
        If reset_for_eval is set to False, then each eval loop will process this
        many batches. Metrics over those batches will be aggregated and then
        reported.
      is_training: Whether or not this dataset is used for model traning.
      experimental_remote_input: whether to process inputs on remote hosts, when
        there is a single controller.
      batch_padding_size: The amount of right-padding applied to each invocation
        of get_next_padded(). Useful when the batch_size is smaller than the
        number of devices.
    """
    batch_size: Optional[int] = None
    # Sharding behavior. If num_infeed_hosts is 0, it will be given a default
    # value by PAX trainer; if it is still not set during __init__, 1 will be
    # used.
    num_infeed_hosts: int = 0
    infeed_host_index: int = 0
    # Deterministic randomness.
    input_random_seed: Optional[int] = None
    reset_for_eval: bool = False
    eval_loop_num_batches: int = 1
    is_training: bool = False
    experimental_remote_input: bool = False
    batch_padding_size: int = 0

  @classmethod
  def get_batch_size(cls, hparams: BaseInput.HParams) -> int:
    assert hparams.batch_size is not None
    return hparams.batch_size

  def __init__(self, hparams: BaseInput.HParams) -> None:
    if self._VALIDATE_BATCH_SIZE_NOT_NONE and (hparams.batch_size is None):
      raise ValueError('Must specify p.batch_size.')
    if not hparams.name:
      hparams.name = 'train' if hparams.is_training else 'input'
    if hparams.num_infeed_hosts == 0:
      hparams.num_infeed_hosts = 1
    super().__init__(hparams)
    name = self.hparams.name
    if re.fullmatch(r'[a-zA-Z0-9.:_-]+', name) is None:
      raise ValueError(f'Input hparams p.name string invalid: "{name}" '
                       'does not fully match "[a-zA-Z0-9._-]+".')
    if hparams.experimental_remote_input and jax.process_count() > 1:
      raise NotImplementedError(
          'Remote input is not supported when there are multiple controllers.')

  def get_next(self) -> NestedJTensor:
    raise NotImplementedError

  def get_next_padded(self) -> NestedJTensor:
    unpadded = self.get_next()
    pad_size = self.hparams.batch_padding_size
    if pad_size == 0:
      return unpadded
    return tf.nest.map_structure(
        lambda x: np.pad(x, [[0, pad_size]] + [[0, 0]] * (x.ndim - 1)),
        unpadded)

  def reset(self) -> None:
    pass

  def ids_to_strings(self,
                     ids: pytypes.NpTensor,
                     lengths: pytypes.NpTensor,
                     key: Optional[str] = None) -> Sequence[str]:
    """Converts int ids into strings.

    Args:
      ids: A matrix of shape [batch, seqlen], each row is a sequence to be
        converted.
      lengths: A vector of shape [batch]. lens[i] is the sequence length of the
        i-th row. Only the first lens[i] tokens in ids[i, :] are valid tokens.
      key: Optional argument to specify which tokenizer to use. This is useful
        for example in a sequence model where the source and targets have
        different tokenizers.

    Returns:
      A list strings of shape [batch]. The converted texts.
    """
    raise NotImplementedError


class LingvoInputAdaptor(BaseInput):
  """Syntactic sugar for adapting a Lingvo style input for Jax.

  This should be able to wrap any Lingvo TF input generator to be used in
  Lingvo Jax. Remember to set `p.is_training=True` on the training dataset.

  Some usage caveats below.

  For eval, `p.num_samples` or other similar params like samples_per_summary are
  completely ignored by Lingvo Jax. Caller should instead set `p.num_batches` to
  (p.num_samples // batch_size) with `p.reset_for_eval=True` so that each eval
  step reads (approximately) one epoch of eval data. This might not be needed if
  the input already is finite (e.g. with p.repeat_count=1).

  When multiple infeed hosts are used, one must take care to ensure that the
  Lingvo input either already uses InfeedContextScope for proper sharding, or
  alternatively do not use the same random seed on all hosts. In other words,
  one must avoid the failure case where each host emits identical training data.
  See also p.allow_fixed_file_random_seed below.
  """
  _VALIDATE_BATCH_SIZE_NOT_NONE = False
  _VALIDATE_BATCH_SIZE_NONE = True<commit_msg>Allow input implementations to provide custom input resharding implementations<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.

"""Base classes for the Praxis input layers."""

from __future__ import annotations

import abc
import copy
import re
from typing import Any, Callable, Dict, Iterable, Optional, Sequence

from absl import logging
import jax
from jax.experimental import maps
from lingvo.core import cluster_factory
from lingvo.core import datasource
import numpy as np
from praxis import base_hyperparams
from praxis import py_utils
from praxis import pytypes
import tensorflow.compat.v2 as tf

NestedMap = py_utils.NestedMap
NestedJTensor = pytypes.NestedJTensor
Nested = pytypes.Nested
NestedShapeDtypeStruct = pytypes.NestedShapeDtypeStruct
instantiate = base_hyperparams.instantiate


class BaseInput(base_hyperparams.BaseParameterizable):
  """Base class for Jax input classes.

  During paxml's train, on each host an input instance will be
  created (instantiate(input_p)), and then get_next() is iteratively
  called in eager mode to generate one batch of data for each step
  of train/eval/etc.

  If supported, for eval, reset() is called after each eval step.
  See p.reset_for_eval below.

  A tf.data based input should inherit this class directly and implement
  get_next() and reset(). For an example of how to handle sharding for both
  training and eval data, please refer to the implementation of
  TFRecordBertInput at tasks/lm/input_generator.py.

  If there is already an Lingvo TF input generator that one would like to
  use directly, please use LingvoInputAdaptor below.
  """
  _VALIDATE_BATCH_SIZE_NOT_NONE = True

  class HParams(base_hyperparams.BaseParameterizable.HParams):
    """Hyper-parameters associated with this Input class.

    Attributes:
      batch_size: The (Jax per process) Batch size. Each call to get_next()
        returns a batch with this batch size.
      num_infeed_hosts: Usually set to jax.process_count(). Implementation must
        ensure that the data is sharded into this many shard.
      infeed_host_index: Usually set to jax.process_index(). Implementation must
        ensure that each instance returns a shard with this index.
      input_random_seed: If set, implementation must ensure that this is used to
        seed randomness, e.g. when shuffling in a deterministic manner.
      reset_for_eval: If set, eval will continue until tf.errors.OutOfRange is
        raised, and reset() will called for each eval. Implementation must
        ensure that all variant p.infeed_host_index instances raise after the
        same numberof calls to get_next() to ensure synchronization across
        hosts. If notset, get_next() must never raise.
      eval_loop_num_batches: Num of batches to process per eval loop. Must be >=
        1. This valueis ignored if reset_for_eval is set True, in which case,
        this value is dynamically determined by the number of available batches.
        If reset_for_eval is set to False, then each eval loop will process this
        many batches. Metrics over those batches will be aggregated and then
        reported.
      is_training: Whether or not this dataset is used for model traning.
      experimental_remote_input: whether to process inputs on remote hosts, when
        there is a single controller.
      batch_padding_size: The amount of right-padding applied to each invocation
        of get_next_padded(). Useful when the batch_size is smaller than the
        number of devices.
    """
    batch_size: Optional[int] = None
    # Sharding behavior. If num_infeed_hosts is 0, it will be given a default
    # value by PAX trainer; if it is still not set during __init__, 1 will be
    # used.
    num_infeed_hosts: int = 0
    infeed_host_index: int = 0
    # Deterministic randomness.
    input_random_seed: Optional[int] = None
    reset_for_eval: bool = False
    eval_loop_num_batches: int = 1
    is_training: bool = False
    experimental_remote_input: bool = False
    batch_padding_size: int = 0

  @classmethod
  def get_batch_size(cls, hparams: BaseInput.HParams) -> int:
    assert hparams.batch_size is not None
    return hparams.batch_size

  def __init__(self, hparams: BaseInput.HParams) -> None:
    if self._VALIDATE_BATCH_SIZE_NOT_NONE and (hparams.batch_size is None):
      raise ValueError('Must specify p.batch_size.')
    if not hparams.name:
      hparams.name = 'train' if hparams.is_training else 'input'
    if hparams.num_infeed_hosts == 0:
      hparams.num_infeed_hosts = 1
    super().__init__(hparams)
    name = self.hparams.name
    if re.fullmatch(r'[a-zA-Z0-9.:_-]+', name) is None:
      raise ValueError(f'Input hparams p.name string invalid: "{name}" '
                       'does not fully match "[a-zA-Z0-9._-]+".')
    if hparams.experimental_remote_input and jax.process_count() > 1:
      raise NotImplementedError(
          'Remote input is not supported when there are multiple controllers.')

  def get_next(self) -> NestedJTensor:
    raise NotImplementedError

  def get_next_padded(self) -> NestedJTensor:
    unpadded = self.get_next()
    pad_size = self.hparams.batch_padding_size
    if pad_size == 0:
      return unpadded
    return tf.nest.map_structure(
        lambda x: np.pad(x, [[0, pad_size]] + [[0, 0]] * (x.ndim - 1)),
        unpadded)

  def reset(self) -> None:
    pass

  def ids_to_strings(self,
                     ids: pytypes.NpTensor,
                     lengths: pytypes.NpTensor,
                     key: Optional[str] = None) -> Sequence[str]:
    """Converts int ids into strings.

    Args:
      ids: A matrix of shape [batch, seqlen], each row is a sequence to be
        converted.
      lengths: A vector of shape [batch]. lens[i] is the sequence length of the
        i-th row. Only the first lens[i] tokens in ids[i, :] are valid tokens.
      key: Optional argument to specify which tokenizer to use. This is useful
        for example in a sequence model where the source and targets have
        different tokenizers.

    Returns:
      A list strings of shape [batch]. The converted texts.
    """
    raise NotImplementedError

  @classmethod
  def reshard_for_pmap(cls, arrays: NestedJTensor) -> NestedJTensor:
    """Reshards inputs for pmap.

    This function reshards `arrays`, inputs returned by this input class, to be
    suitable for pmap computation. The returned arrays are expected to have the
    pmap device dimension as the 0th dimension and per-device batch dimension as
    the 1st dimension. Input subclasses may override this to customize the
    resharding implementations.

    Args:
      arrays: Inputs returned by this input class.

    Returns:
      Resharded inputs.
    """
    return jax.tree_util.tree_map(py_utils.reshard, arrays)

  @classmethod
  def reshard_for_spmd(cls, arrays: NestedJTensor,
                       global_shapes: NestedShapeDtypeStruct,
                       global_mesh: maps.Mesh, pspecs: Any) -> NestedJTensor:
    """Reshards inputs for pjit.

    This function reshards `arrays`, inputs returned by this input class, to be
    suitable for pjit computation on the given mesh. The caller also provides
    the expected global shapes and partition specs of the arrays to be returned.
    The returned arrays are expected to have the global batch dimension as the
    0th dimension and be sharded in the way specified by `pspecs`.Input
    subclasses may override this to customize the resharding implementations.

    Args:
      arrays: Inputs returned by this input class.
      global_shapes: Expected global shapes of `arrays` after resharding.
      global_mesh: Global mesh for pjit computation.
      pspecs: Expected partition specs for `arrays` after resharding.

    Returns:
      Resharded inputs.
    """
    py_utils.assert_same_shape_and_dtype(
        global_shapes,
        tf.nest.map_structure(py_utils.get_global_input_shape_dtype, arrays))
    return py_utils.create_gda(arrays, global_shapes, global_mesh, pspecs)


class LingvoInputAdaptor(BaseInput):
  """Syntactic sugar for adapting a Lingvo style input for Jax.

  This should be able to wrap any Lingvo TF input generator to be used in
  Lingvo Jax. Remember to set `p.is_training=True` on the training dataset.

  Some usage caveats below.

  For eval, `p.num_samples` or other similar params like samples_per_summary are
  completely ignored by Lingvo Jax. Caller should instead set `p.num_batches` to
  (p.num_samples // batch_size) with `p.reset_for_eval=True` so that each eval
  step reads (approximately) one epoch of eval data. This might not be needed if
  the input already is finite (e.g. with p.repeat_count=1).

  When multiple infeed hosts are used, one must take care to ensure that the
  Lingvo input either already uses InfeedContextScope for proper sharding, or
  alternatively do not use the same random seed on all hosts. In other words,
  one must avoid the failure case where each host emits identical training data.
  See also p.allow_fixed_file_random_seed below.
  """
  _VALIDATE_BATCH_SIZE_NOT_NONE = False
  _VALIDATE_BATCH_SIZE_NONE = True
{'fraction_non_alphanumeric': 0.06134405514072372, 'fraction_numerical': 0.0036186099942561744, 'mean_word_length': 3.369134253450439, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 21, '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': '6571830', 'n_tokens_mistral': 5025, 'n_tokens_neox': 4850, 'n_words': 2038}