Method load

official/nlp/data/data_loader.py:27–48

Implements DataLoader load method. Builds the entire input pipeline inside the load method. Users can define states inside the DataLoader class and returns a tf.data dataset object. Args: input_context: This is a context class that is passed to the user's input functi

(
      self,
      input_context: Optional[tf.distribute.InputContext] = None
  )

  @abc.abstractmethod
  def load(
      self,
      input_context: Optional[tf.distribute.InputContext] = None
  ) -> tf.data.Dataset:
    """Implements DataLoader load method.

    Builds the entire input pipeline inside the load method. Users can define
    states inside the DataLoader class and returns a tf.data dataset
    object.

    Args:
      input_context: This is a context class that is passed to the user's input
        function and contains information about the compute replicas and input
        pipelines. This object is used for multi-host inputs and passed by the
        distribution strategy.

    Returns:
      A per-host tf.data dataset. Note that, we usually create the distributed
        dataset through the load method, so we should not directly return a
        distributed dataset here.
    """
    pass