def consume_noise_mask(vocab_size=32100):
  """Consumes (tokens, noise mask) and returns (inputs, targets)."""
  def _noise_span_to_unique_sentinel(tokens, noise_mask):
    prev_token_is_noise = np.pad(
        noise_mask[:-1], [1, 0], mode='constant', constant_values=False)
    first_noise_tokens = np.logical_and(noise_mask,
                                        np.logical_not(prev_token_is_noise))
    subsequent_noise_tokens = np.logical_and(noise_mask, prev_token_is_noise)
    sentinel = vocab_size - np.cumsum(first_noise_tokens)
    tokens = np.where(first_noise_tokens, sentinel, tokens)
    return tokens[np.logical_not(subsequent_noise_tokens)]

  def _f(generator):
    for tokens, noise_mask in generator:
      # Returns inputs and targets.
      yield (_noise_span_to_unique_sentinel(tokens, noise_mask),
             _noise_span_to_unique_sentinel(tokens, np.logical_not(noise_mask)))
  return _f
Consumes (tokens, noise mask) and returns (inputs, targets).
trax/data/inputs.py