group_split(
decisions: Sequence[Decision],
*,
eval_fraction: float = 0.2,
seed: int = 0
) -> tuple[list[Decision], list[Decision]]
Split into (train, eval) on Decision.grouping.
Deterministic given seed: groups are sorted by key before being shuffled, so the
result does not depend on the order the decisions were loaded in. Rows keep their
input order within each side.
Raises DataError if the eval side would hold zero examples of some (criterion,
answer) class, because a class absent from eval is a class the score says nothing
about.
Source code in src/typedecide/data/split.py
| def group_split(
decisions: Sequence[Decision],
*,
eval_fraction: float = 0.2,
seed: int = 0,
) -> tuple[list[Decision], list[Decision]]:
"""Split into (train, eval) on `Decision.grouping`.
Deterministic given `seed`: groups are sorted by key before being shuffled, so the
result does not depend on the order the decisions were loaded in. Rows keep their
input order within each side.
Raises `DataError` if the eval side would hold zero examples of some (criterion,
answer) class, because a class absent from eval is a class the score says nothing
about.
"""
if not decisions:
raise DataError(
"Cannot split an empty dataset. Fix: check the loader's source and any filter "
"applied after loading."
)
if not 0.0 < eval_fraction < 1.0:
raise DataError(
f"eval_fraction={eval_fraction} must be strictly between 0 and 1. "
"Fix: pass something like 0.2 for a 20% eval side."
)
by_group: dict[str, list[Decision]] = defaultdict(list)
for decision in decisions:
by_group[decision.grouping].append(decision)
keys = sorted(by_group)
if len(keys) < 2:
raise DataError(
f"All {len(decisions)} rows share the group {keys[0]!r}, so a group split would put "
"every row on one side. Fix: set group_id to the thing that must not be split -- "
"usually the state or the source ticket -- or leave it unset so each row is its own "
"group."
)
rng = random.Random(seed)
shuffled = list(keys)
rng.shuffle(shuffled)
total = len(decisions)
target = eval_fraction * total
eval_keys: list[str] = []
taken = 0
for key in shuffled:
if taken >= target or len(eval_keys) == len(shuffled) - 1:
break
eval_keys.append(key)
taken += len(by_group[key])
if not eval_keys:
eval_keys = [shuffled[0]]
held = set(eval_keys)
_refuse_if_class_missing(decisions, by_group, held, eval_fraction, seed)
train = [d for d in decisions if d.grouping not in held]
evaluation = [d for d in decisions if d.grouping in held]
LOGGER.info(
"split %d rows in %d groups into %d train / %d eval (seed=%d)",
total,
len(keys),
len(train),
len(evaluation),
seed,
)
return train, evaluation
|