def METHOD_NAME(self, num_elements, batch_size, k):
    """Test the python implementation."""
    ds = tf.data.Dataset.from_generator(
        word_gen, tf.string, tf.TensorShape([])
    )
    batched_ds = ds.take(num_elements).batch(batch_size)
    input_t = keras.Input(shape=(), dtype=tf.string)
    layer = index_lookup.IndexLookup(
        max_tokens=k,
        num_oov_indices=0,
        mask_token=None,
        oov_token="OOV",
        dtype=tf.string,
    )
    _ = layer(input_t)
    num_repeats = 5
    starts = []
    ends = []
    for _ in range(num_repeats):
        starts.append(time.time())
        vocab = get_top_k(batched_ds, k)
        layer.set_vocabulary(vocab)
        ends.append(time.time())
    avg_time = np.mean(np.array(ends) - np.array(starts))
    return avg_time