gpu/single-gpu/home_benjamin_hdbasin/code/hdbasin/real_data.py

from __future__ import annotations

import time
from dataclasses import dataclass


@dataclass
class RealDatasetSpec:
    task: str
    dataset_id: str
    config: str | None
    objective: str


REAL_DATASETS = [
    RealDatasetSpec("Fashion-MNIST", "zalando-datasets/fashion_mnist", None, "Small CNN validation cross-entropy"),
    RealDatasetSpec("CIFAR-10", "uoft-cs/cifar10", None, "Tiny CNN or tiny ResNet validation loss"),
    RealDatasetSpec("AG News", "fancyzhx/ag_news", None, "TF-IDF + logistic regression validation loss"),
    RealDatasetSpec("SST-2", "nyu-mll/glue", "sst2", "DistilBERT fine-tuning validation loss"),
    RealDatasetSpec("IMDB", "stanfordnlp/imdb", None, "DistilBERT fine-tuning validation loss"),
    RealDatasetSpec("Tabular", "inria-soda/tabular-benchmark", None, "Tree/MLP validation loss by task"),
]


def probe_real_datasets() -> list[dict]:
    rows = []
    try:
        import torch
        from datasets import get_dataset_config_names, load_dataset_builder
    except Exception as exc:
        now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        return [
            {
                "task": spec.task,
                "dataset_id": spec.dataset_id,
                "config": spec.config or "",
                "objective": spec.objective,
                "status": "not_run",
                "reason": f"required package unavailable: {type(exc).__name__}: {exc}",
                "checked_at_utc": now,
                "gpu_available": "",
                "gpu_name": "",
            }
            for spec in REAL_DATASETS
        ]
    gpu_available = bool(torch.cuda.is_available())
    gpu_name = torch.cuda.get_device_name(0) if gpu_available else ""
    now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    for spec in REAL_DATASETS:
        try:
            if spec.config is None:
                # This fetches builder metadata, not the full dataset.
                builder = load_dataset_builder(spec.dataset_id)
            else:
                builder = load_dataset_builder(spec.dataset_id, spec.config)
            rows.append(
                {
                    "task": spec.task,
                    "dataset_id": spec.dataset_id,
                    "config": spec.config or "",
                    "objective": spec.objective,
                    "status": "available_metadata_only",
                    "reason": f"builder loaded; splits={','.join(builder.info.splits.keys()) if builder.info.splits else 'unknown'}",
                    "checked_at_utc": now,
                    "gpu_available": gpu_available,
                    "gpu_name": gpu_name,
                }
            )
        except Exception as exc:
            # Some HF dataset repos require configs; record instead of silently skipping.
            try:
                configs = get_dataset_config_names(spec.dataset_id)
                reason = f"builder failed; available configs sample={configs[:8]}; {type(exc).__name__}: {exc}"
            except Exception:
                reason = f"builder failed; {type(exc).__name__}: {exc}"
            rows.append(
                {
                    "task": spec.task,
                    "dataset_id": spec.dataset_id,
                    "config": spec.config or "",
                    "objective": spec.objective,
                    "status": "not_run",
                    "reason": reason,
                    "checked_at_utc": now,
                    "gpu_available": gpu_available,
                    "gpu_name": gpu_name,
                }
            )
    return rows