<commit_before>import numpy        as np
import numpy.random as npr

import kayak

from . import *
<commit_msg>Some tests for the batcher.<commit_after>import numpy        as np
import numpy.random as npr

import kayak

from . import *

def test_indices_1():
    """Test with deterministic indices."""

    for num_data in [1, 10, 100, 1000, 10000, 100000, 1000000]:
        for batch_size in [1, 10, 11, 25, 50, 101, 500, 1000, 1011]:
            
            data_used = np.zeros((num_data,), dtype=bool)
            batcher = kayak.Batcher(batch_size, num_data)
            for batch in batcher:
                data_used[batch] = True
            
            assert np.all(data_used)

def test_indices_2():
    """Test with random seed."""
    npr.seed(1)

    for num_data in [1, 10, 100, 1000, 10000, 100000, 1000000]:
        for batch_size in [1, 10, 11, 25, 50, 101, 500, 1000, 1011]:
            
            data_used = np.zeros((num_data,), dtype=bool)
            batcher = kayak.Batcher(batch_size, num_data, rng=npr.randint(1,100000))
            for batch in batcher:
                data_used[batch] = True
            
            assert np.all(data_used)

def test_indices_3():
    """Test with RandomState object."""
    npr.seed(2)

    for num_data in [1, 10, 100, 1000, 10000, 100000, 1000000]:
        for batch_size in [1, 10, 11, 25, 50, 101, 500, 1000, 1011]:
            
            data_used = np.zeros((num_data,), dtype=bool)
            batcher = kayak.Batcher(batch_size, num_data, rng=npr.RandomState(npr.randint(1,100000)))
            for batch in batcher:
                data_used[batch] = True
            
            assert np.all(data_used)

def test_reset():
    """Test resetting."""

    for num_data in [1000, 10000, 100000, 1000000]:
        for batch_size in [1, 10, 11, 25, 50, 101, 500]:
            
            batcher = kayak.Batcher(batch_size, num_data)

            # Start the batcher forward.
            batcher.next()

            # Now reset.
            batcher.reset()

            # Make sure we touch all of the data.
            data_used = np.zeros((num_data,), dtype=bool)
            for batch in batcher:
                data_used[batch] = True
            
            assert np.all(data_used)
{'fraction_non_alphanumeric': 0.10812003530450133, 'fraction_numerical': 0.09223300970873786, 'mean_word_length': 1.9064102564102565, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5286725', 'n_tokens_mistral': 908, 'n_tokens_neox': 700, 'n_words': 209}