datasets

1
import json
2
import os
3
import tempfile
4

5
import datasets
6
from utils import generate_example_dataset, get_duration
7

8

9
SPEED_TEST_N_EXAMPLES = 500_000
10

11
RESULTS_BASEPATH, RESULTS_FILENAME = os.path.split(__file__)
12
RESULTS_FILE_PATH = os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))
13

14

15
@get_duration
16
def select(dataset: datasets.Dataset):
17
    _ = dataset.select(range(0, len(dataset), 2))
18

19

20
@get_duration
21
def sort(dataset: datasets.Dataset):
22
    _ = dataset.sort("numbers")
23

24

25
@get_duration
26
def shuffle(dataset: datasets.Dataset):
27
    _ = dataset.shuffle()
28

29

30
@get_duration
31
def train_test_split(dataset: datasets.Dataset):
32
    _ = dataset.train_test_split(0.1)
33

34

35
@get_duration
36
def shard(dataset: datasets.Dataset, num_shards=10):
37
    for shard_id in range(num_shards):
38
        _ = dataset.shard(num_shards, shard_id)
39

40

41
def benchmark_indices_mapping():
42
    times = {"num examples": SPEED_TEST_N_EXAMPLES}
43
    functions = (select, sort, shuffle, train_test_split, shard)
44
    with tempfile.TemporaryDirectory() as tmp_dir:
45
        print("generating dataset")
46
        features = datasets.Features({"text": datasets.Value("string"), "numbers": datasets.Value("float32")})
47
        dataset = generate_example_dataset(
48
            os.path.join(tmp_dir, "dataset.arrow"), features, num_examples=SPEED_TEST_N_EXAMPLES
49
        )
50
        print("Functions")
51
        for func in functions:
52
            print(func.__name__)
53
            times[func.__name__] = func(dataset)
54

55
    with open(RESULTS_FILE_PATH, "wb") as f:
56
        f.write(json.dumps(times).encode("utf-8"))
57

58

59
if __name__ == "__main__":  # useful to run the profiler
60
    benchmark_indices_mapping()
61