Spaces:
Paused
Paused
File size: 1,258 Bytes
0d20347 | 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 | """Leakage-resistant deterministic train/val/test splitting."""
from __future__ import annotations
from hashlib import blake2b
from typing import Iterable
from .dedup import normalized_hash
from .document import CorpusDocument
def split_documents(
docs: Iterable[CorpusDocument],
*,
train_ratio: float = 0.9,
val_ratio: float = 0.05,
test_ratio: float = 0.05,
seed: str = "cogcore-v014",
) -> list[CorpusDocument]:
total = train_ratio + val_ratio + test_ratio
if not 0.999 <= total <= 1.001:
raise ValueError("train_ratio + val_ratio + test_ratio must be 1.0")
out: list[CorpusDocument] = []
assigned_by_hash: dict[str, str] = {}
for doc in docs:
group = normalized_hash(doc.text)
split = assigned_by_hash.get(group)
if split is None:
key = f"{seed}:{group}".encode("utf-8")
val = int.from_bytes(blake2b(key, digest_size=8).digest(), "big") / float(2**64 - 1)
if val < train_ratio:
split = "train"
elif val < train_ratio + val_ratio:
split = "val"
else:
split = "test"
assigned_by_hash[group] = split
out.append(doc.with_split(split))
return out
|