def generate_sequential_chunks(max_length=None):
  """Returns a function that generates chunks of atmost max_length length."""
  def _f(generator):
    for example in generator:
      n_tokens = len(example)
      if n_tokens <= max_length:
        yield example
      else:
        n_segments = int(math.ceil(float(n_tokens) / float(max_length)))
        for i in range(n_segments):
          start = max_length * i
          end = min(start + max_length, n_tokens)
          yield example[start:end]
  return _f
Returns a function that generates chunks of atmost max_length length.
trax/data/inputs.py