cogcore/data/split.py

"""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