Function stratified_3way_split

scripts/split_data.py:17–47

(
    rows: list[dict[str, Any]],
    seed: int = 42,
    train_frac: float = 0.64,
    cal_frac: float = 0.16,
)

def stratified_3way_split(
    rows: list[dict[str, Any]],
    seed: int = 42,
    train_frac: float = 0.64,
    cal_frac: float = 0.16,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]], list[dict[str, Any]]]:
    rng = random.Random(seed)
    strata: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        key = f"{row.get('benchmark', 'unknown')}_{row.get('target_tier', 'unknown')}"
        strata[key].append(row)

    train, cal, holdout = [], [], []
    for key in sorted(strata.keys()):
        group = list(strata[key])
        rng.shuffle(group)
        n = len(group)
        n_train = max(1, round(n * train_frac))
        n_cal = max(0, round(n * cal_frac))
        if n > 1:
            n_holdout = n - n_train - n_cal
            if n_holdout < 1:
                n_cal = max(0, n_cal - 1)
        train.extend(group[:n_train])
        cal.extend(group[n_train:n_train + n_cal])
        holdout.extend(group[n_train + n_cal:])

    rng.shuffle(train)
    rng.shuffle(cal)
    rng.shuffle(holdout)
    return train, cal, holdout


def main():