def advance_iter_and_group_samples(train_iterator, num_samples, max_seq_length):
    """
    The training iterator is advanced so that after groupifying the samples,
    `num_samples` of length `max_seq_length` are returned.
    """
    num_total_tokens = max_seq_length * num_samples
    samples = defaultdict(list)

    i = 0
    while i < num_total_tokens:
        tokenized_samples = next(train_iterator)
        i += len(tokenized_samples["input_ids"])

        # concatenate tokenized samples to list (excluding "id" and "text")
        samples = {
            k: samples[k] + tokenized_samples[k] for k in ["input_ids", "attention_mask", "special_tokens_mask"]
        }

    # Concatenated tokens are split to lists of length `max_seq_length`.
    # Note that remainedr of % max_seq_length are thrown away.
    def group_texts(examples):
        result = {
            k: [t[i : i + max_seq_length] for i in range(0, num_total_tokens, max_seq_length)]
            for k, t in examples.items()
        }
        return result

    grouped_samples = group_texts(samples)
    return grouped_samples

    The training iterator is advanced so that after groupifying the samples,
    `num_samples` of length `max_seq_length` are returned.
    
dpkd/transformers/examples/research_projects/jax-projects/dataset-streaming/run_mlm_flax_stream.py