def __init__(self, data_stream, vocab_size, metric=None,
               eos_id=1, max_length=1000):
    """The constructor.

    Args:
      data_stream: A python generator creating lists or tuples of
        sequences (list, tuples or numpy arrays) of integers.
      vocab_size: Integer, the size of the vocabulary. All integers in the
        data stream must be positive and smaller than this value.
      metric: A function taking two lists of integers and returning a float.
        If None, we use per-token accuracy as the default metric.
      eos_id: Integer, the id of the EOS symbol.
      max_length: Integer, maximum length of the policy reply to avoid
        infinite episodes if policy never produces EOS.

    Returns:
      A new environment which presents the data and compares the policy
      response with the expected data, returning metric as reward.
    """
    self._data = data_stream
    self._vocab_size = vocab_size
    self._eos = eos_id
    self._max_length = max_length
    self._metric = _accuracy if metric is None else metric
    self.reset()
The constructor.

    Args:
      data_stream: A python generator creating lists or tuples of
        sequences (list, tuples or numpy arrays) of integers.
      vocab_size: Integer, the size of the vocabulary. All integers in the
        data stream must be positive and smaller than this value.
      metric: A function taking two lists of integers and returning a float.
        If None, we use per-token accuracy as the default metric.
      eos_id: Integer, the id of the EOS symbol.
      max_length: Integer, maximum length of the policy reply to avoid
        infinite episodes if policy never produces EOS.

    Returns:
      A new environment which presents the data and compares the policy
      response with the expected data, returning metric as reward.
    
trax/rl/envs/data_envs.py