File size: 1,731 Bytes
ee50dca | 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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | from __future__ import annotations
import json
from pathlib import Path
import numpy as np
from datasets import Dataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
def split_indices(labels: np.ndarray, seed: int = 2026) -> dict[str, np.ndarray]:
indices = np.arange(len(labels))
train, remainder = train_test_split(
indices,
test_size=0.30,
stratify=labels,
random_state=seed,
)
validation, test = train_test_split(
remainder,
test_size=0.50,
stratify=labels[remainder],
random_state=seed,
)
return {"train": train, "validation": validation, "test": test}
def main() -> None:
DATA_DIR.mkdir(parents=True, exist_ok=True)
digits = load_digits()
images = digits.images.astype(np.float32)
labels = digits.target.astype(np.int64)
manifest = {}
for name, indices in split_indices(labels).items():
dataset = Dataset.from_dict(
{
"image": [image.reshape(-1).tolist() for image in images[indices]],
"label": labels[indices].tolist(),
}
)
path = DATA_DIR / f"{name}.parquet"
dataset.to_parquet(path)
counts = np.bincount(labels[indices], minlength=10)
manifest[name] = {
"rows": len(indices),
"class_counts": counts.tolist(),
"path": path.name,
}
(DATA_DIR / "manifest.json").write_text(
json.dumps(manifest, indent=2),
encoding="utf-8",
)
print(json.dumps(manifest, indent=2))
if __name__ == "__main__":
main()
|