File size: 2,884 Bytes
2dd5f57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """The train and val splits must be exactly complementary for any train_ratio."""
from pathlib import Path
from typing import Literal
import pytest
from datasets import Dataset
from speculators.train.data import ArrowDataset
def _dataset(tmp_path: Path, n: int) -> str:
ds = Dataset.from_dict(
{
"input_ids": [[i, i + 1, i + 2] for i in range(n)],
"loss_mask": [[1, 1, 1]] * n,
"seq_len": [3] * n,
}
)
path = tmp_path / "data"
ds.save_to_disk(str(path))
(path / "hidden_states").mkdir()
return str(path)
def _split(path: str, ratio: float, split: Literal["train", "val"]) -> ArrowDataset:
return ArrowDataset(
max_len=128,
datapath=path,
on_missing="skip",
train_ratio=ratio,
split=split,
)
@pytest.mark.parametrize("ratio", [0.1, 0.2, 0.25, 0.3, 0.5, 0.7, 0.8, 0.9, 0.95])
@pytest.mark.parametrize("n", [1000, 12501, 100000])
def test_splits_are_exactly_complementary(tmp_path, ratio, n):
path = _dataset(tmp_path, n)
train = _split(path, ratio, "train")
val = _split(path, ratio, "val")
assert len(train) + len(val) == n, "splits must partition the dataset"
assert len(train) == int(n * ratio)
assert val.start_file_idx == len(train), "val must begin where train ends"
assert len(train) > 0
assert len(val) > 0
@pytest.mark.parametrize("ratio", [0.2, 0.9])
def test_no_row_appears_in_both_splits(tmp_path, ratio):
n = 12501
path = _dataset(tmp_path, n)
train = _split(path, ratio, "train")
val = _split(path, ratio, "val")
train_ids = {tuple(train.data[i]["input_ids"]) for i in range(len(train))}
val_ids = {tuple(val.data[i]["input_ids"]) for i in range(len(val))}
assert not (train_ids & val_ids), "a row is in both train and val"
def test_default_takes_whole_dataset(tmp_path):
path = _dataset(tmp_path, 100)
assert len(_split(path, 1.0, "train")) == 100
@pytest.mark.parametrize("ratio", [0.0, -0.1, 1.5])
def test_rejects_out_of_range_ratio(tmp_path, ratio):
path = _dataset(tmp_path, 100)
with pytest.raises(ValueError, match="train_ratio must be in"):
_split(path, ratio, "train")
def test_rejects_val_split_with_no_val_data(tmp_path):
path = _dataset(tmp_path, 100)
with pytest.raises(ValueError, match="leaves no validation split"):
_split(path, 1.0, "val")
@pytest.mark.parametrize("ratio", [0.1, 0.5, 0.9])
def test_rejects_empty_train_from_small_dataset(tmp_path, ratio):
path = _dataset(tmp_path, 1)
with pytest.raises(ValueError, match="train split is empty"):
_split(path, ratio, "train")
def test_small_dataset_both_splits_nonempty(tmp_path):
path = _dataset(tmp_path, 2)
train = _split(path, 0.5, "train")
val = _split(path, 0.5, "val")
assert len(train) == 1
assert len(val) == 1
|