Spaces:
Runtime error
Runtime error
File size: 757 Bytes
f8f62f3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
import pytest
from training.data import get_dataset_size
@pytest.mark.parametrize(
"shards,expected_size",
[
('/path/to/shard.tar', 1),
('/path/to/shard_{000..000}.tar', 1),
('/path/to/shard_{000..009}.tar', 10),
('/path/to/shard_{000..009}_{000..009}.tar', 100),
('/path/to/shard.tar::/path/to/other_shard_{000..009}.tar', 11),
('/path/to/shard_{000..009}.tar::/path/to/other_shard_{000..009}.tar', 20),
(['/path/to/shard.tar'], 1),
(['/path/to/shard.tar', '/path/to/other_shard.tar'], 2),
]
)
def test_num_shards(shards, expected_size):
_, size = get_dataset_size(shards)
assert size == expected_size, f'Expected {expected_size} for {shards} but found {size} instead.'
|