examples/tensorflow/translation/run_translation.py/sample_generator
def sample_generator(dataset, model, tokenizer, shuffle, pad_to_multiple_of=None):
    if shuffle:
        sample_ordering = np.random.permutation(len(dataset))
    else:
        sample_ordering = np.arange(len(dataset))
    for sample_idx in sample_ordering:
        example = dataset[int(sample_idx)]
        # Handle dicts with proper padding and conversion to tensor.
        example = tokenizer.pad(example, return_tensors="np", pad_to_multiple_of=pad_to_multiple_of)
        example = {key: tf.convert_to_tensor(arr, dtype_hint=tf.int32) for key, arr in example.items()}
        if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"):
            decoder_input_ids = model.prepare_decoder_input_ids_from_labels(
                labels=tf.expand_dims(example["labels"], 0)
            )
            example["decoder_input_ids"] = tf.squeeze(decoder_input_ids, 0)
        yield example, example["labels"]  # TF needs some kind of labels, even if we don't use them
    return