def PrefixLM(input_length=128, output_length=512):  # pylint:disable=invalid-name
  """Chunks examples so as to make inputs/outputs of specified lenghts."""
  def _f(generator):
    for example in generator:
      n_tokens = len(example)
      # Iterate:
      # |--------|<---- input_length ---->|<- output_length ->|--------------|
      # ^        ^                        ^                   ^
      # |        |                        |                   |
      # 0        input_begin_idx          input_end_idx       output_end_idx
      input_begin_idx = 0
      # While you can make an input batch, keep going.
      while input_begin_idx + input_length < n_tokens:
        input_end_idx = input_begin_idx + input_length
        output_end_idx = min(input_end_idx + output_length, n_tokens)
        yield (example[input_begin_idx:input_end_idx],
               example[input_end_idx:output_end_idx])
        # Update the indices.
        input_begin_idx = output_end_idx
  return _f
Chunks examples so as to make inputs/outputs of specified lenghts.
trax/data/inputs.py