ludwig-ai--ludwig
wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
82 行
2.5 KiB
Python
82 行
2.5 KiB
Python
"""Dataset splitting utilities.
|
|
|
|
Provides functions for splitting datasets into train/validation/test sets
|
|
using various strategies: random, stratified, fixed column, datetime, hash.
|
|
|
|
Extracted from preprocessing.py for modularity. The main split logic remains
|
|
in ludwig/data/split.py; this module provides additional utilities.
|
|
"""
|
|
|
|
import logging
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_split_indices(
|
|
n_samples: int,
|
|
probabilities: tuple[float, float, float] = (0.7, 0.1, 0.2),
|
|
random_seed: int = 42,
|
|
) -> np.ndarray:
|
|
"""Generate split indices (0=train, 1=validation, 2=test) for a dataset.
|
|
|
|
Args:
|
|
n_samples: Number of samples in the dataset.
|
|
probabilities: (train, val, test) split ratios. Must sum to 1.
|
|
random_seed: Random seed for reproducibility.
|
|
|
|
Returns:
|
|
Array of split indices (0, 1, or 2) for each sample.
|
|
"""
|
|
if abs(sum(probabilities) - 1.0) >= 1e-6:
|
|
raise ValueError(
|
|
f"Split probabilities must sum to 1, got {sum(probabilities)}.\n"
|
|
f"Fix: ensure your train/validation/test split fractions sum to 1.0 (e.g., [0.7, 0.1, 0.2])."
|
|
)
|
|
|
|
rng = np.random.RandomState(random_seed)
|
|
indices = rng.permutation(n_samples)
|
|
splits = np.zeros(n_samples, dtype=int)
|
|
|
|
train_end = int(n_samples * probabilities[0])
|
|
val_end = train_end + int(n_samples * probabilities[1])
|
|
|
|
splits[indices[train_end:val_end]] = 1
|
|
splits[indices[val_end:]] = 2
|
|
|
|
return splits
|
|
|
|
|
|
def stratified_split_indices(
|
|
labels: np.ndarray,
|
|
probabilities: tuple[float, float, float] = (0.7, 0.1, 0.2),
|
|
random_seed: int = 42,
|
|
) -> np.ndarray:
|
|
"""Generate stratified split indices that maintain label distribution.
|
|
|
|
Args:
|
|
labels: Array of class labels for each sample.
|
|
probabilities: (train, val, test) split ratios.
|
|
random_seed: Random seed for reproducibility.
|
|
|
|
Returns:
|
|
Array of split indices (0, 1, or 2) for each sample.
|
|
"""
|
|
rng = np.random.RandomState(random_seed)
|
|
splits = np.zeros(len(labels), dtype=int)
|
|
|
|
unique_labels = np.unique(labels)
|
|
for label in unique_labels:
|
|
label_indices = np.where(labels == label)[0]
|
|
rng.shuffle(label_indices)
|
|
|
|
n = len(label_indices)
|
|
train_end = int(n * probabilities[0])
|
|
val_end = train_end + int(n * probabilities[1])
|
|
|
|
splits[label_indices[train_end:val_end]] = 1
|
|
splits[label_indices[val_end:]] = 2
|
|
|
|
return splits
|