| """Dataset utilities.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| import numpy as np |
| from sklearn.datasets import load_digits |
| from sklearn.model_selection import train_test_split |
|
|
|
|
| @dataclass |
| class SplitData: |
| X_train: np.ndarray |
| y_train: np.ndarray |
| X_val: np.ndarray |
| y_val: np.ndarray |
| X_test: np.ndarray |
| y_test: np.ndarray |
|
|
|
|
| def load_digits_splits(seed: int = 7, test_size: float = 0.20, val_size: float = 0.20) -> SplitData: |
| """Load standard sklearn Digits and create stratified train/val/test splits.""" |
| digits = load_digits() |
| X = digits.images.astype(np.float32) / 16.0 |
| y = digits.target.astype(np.int64) |
| X_trainval, X_test, y_trainval, y_test = train_test_split( |
| X, y, test_size=test_size, random_state=seed, stratify=y |
| ) |
| val_fraction_of_trainval = val_size / (1.0 - test_size) |
| X_train, X_val, y_train, y_val = train_test_split( |
| X_trainval, y_trainval, test_size=val_fraction_of_trainval, random_state=seed, stratify=y_trainval |
| ) |
| return SplitData(X_train, y_train, X_val, y_val, X_test, y_test) |
|
|
