Function load_data

examples/sldbench/data_loader.py:52–111

Unified data loading interface. Loads and processes data from Hugging Face Hub. Each task's dataset is grouped by a 'group' key. The function returns a dictionary mapping each group key to a tuple of (features, labels). - features (X): A numpy array of shape (n_samples, n_features

(
    app_name: str,
    train: bool = True,
)

}

def load_data(
    app_name: str,
    train: bool = True,
) -> Dict[Any, Tuple[np.ndarray, np.ndarray]]:
    """
    Unified data loading interface. Loads and processes data from Hugging Face Hub.

    Each task's dataset is grouped by a 'group' key. The function returns a 
    dictionary mapping each group key to a tuple of (features, labels).
    - features (X): A numpy array of shape (n_samples, n_features).
    - labels (y): A numpy array of shape (n_samples,) or (n_samples, n_targets).

    Args:
        app_name: The name of the task (e.g., 'sft_scaling_law').
        train: If True, load training data; otherwise, load test data.

    Returns:
        A dictionary containing the prepared data, structured by group.
    """
    if app_name not in TASK_SCHEMA_MAP:
        raise ValueError(f"Task '{app_name}' not found in TASK_SCHEMA_MAP. Available tasks: {list(TASK_SCHEMA_MAP.keys())}")

    split = 'train' if train else 'test'
    schema = TASK_SCHEMA_MAP[app_name]
    
    try:
        # Load the specific task dataset from the Hugging Face Hub
        dataset = datasets.load_dataset(HUB_REPO_ID, name=app_name, split=split)
    except Exception as e:
        raise IOError(f"Failed to load dataset '{app_name}' with split '{split}' from '{HUB_REPO_ID}'. Reason: {e}")

    # Ensure target_name is a list for consistent processing
    feature_names = schema["feature_names"]
    target_names = schema["target_name"]
    if not isinstance(target_names, list):
        target_names = [target_names]

    processed_data = {}
    
    # The datasets are partitioned by a 'group' column
    unique_groups = sorted(list(set(dataset['group'])))

    for group_key in unique_groups:
        # Filter the dataset for the current group
        group_data = dataset.filter(lambda example: example['group'] == group_key)
        
        # Extract features (X) and stack them into a single numpy array
        X_list = [np.array(group_data[fname]) for fname in feature_names]
        X = np.stack(X_list, axis=1)
        
        # Extract targets (y)
        y_list = [np.array(group_data[tname]) for tname in target_names]
        y_stacked = np.stack(y_list, axis=1)
        
        # Squeeze the last dimension if there is only one target
        y = y_stacked.squeeze(axis=1) if y_stacked.shape[1] == 1 else y_stacked
        
        processed_data[group_key] = (X, y)