wehub-resource-sync
85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
"""
|
|
Configuration file listing all datasets used in tests.
|
|
|
|
This file is used by the CI pipeline to pre-download datasets before running tests.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TypedDict
|
|
|
|
|
|
class DictKwargs(TypedDict, total=False):
|
|
subset: str
|
|
categories: list[str]
|
|
|
|
|
|
class Dataset(TypedDict):
|
|
name: str
|
|
kwargs: DictKwargs
|
|
|
|
|
|
# Shap datasets that download from URLs
|
|
SHAP_DATASETS = [
|
|
"imagenet50",
|
|
"california",
|
|
"imdb",
|
|
"adult",
|
|
"nhanesi",
|
|
"a1a",
|
|
"rank",
|
|
"linnerud",
|
|
"diabetes",
|
|
"iris",
|
|
]
|
|
|
|
# Sklearn datasets that need to be fetched (download from internet)
|
|
SKLEARN_FETCH_DATASETS: list[Dataset] = [
|
|
{
|
|
"name": "fetch_california_housing",
|
|
"kwargs": {},
|
|
},
|
|
{
|
|
"name": "fetch_20newsgroups",
|
|
"kwargs": {"subset": "train", "categories": ["alt.atheism", "talk.religion.misc"]},
|
|
},
|
|
{
|
|
"name": "fetch_20newsgroups",
|
|
"kwargs": {"subset": "test", "categories": ["alt.atheism", "talk.religion.misc"]},
|
|
},
|
|
]
|
|
|
|
|
|
def download_all_datasets():
|
|
"""Download all datasets used in tests."""
|
|
from sklearn import datasets as sklearn_datasets
|
|
|
|
import shap
|
|
|
|
print("Downloading shap datasets...")
|
|
for dataset_name in SHAP_DATASETS:
|
|
try:
|
|
dataset_func = getattr(shap.datasets, dataset_name)
|
|
print(f" - {dataset_name}...", end=" ")
|
|
dataset_func()
|
|
print("✓")
|
|
except Exception as e:
|
|
print(f"✗ (Error: {e})")
|
|
|
|
print("\nFetching sklearn datasets...")
|
|
for dataset_config in SKLEARN_FETCH_DATASETS:
|
|
try:
|
|
dataset_func = getattr(sklearn_datasets, dataset_config["name"])
|
|
kwargs_str = ", ".join(f"{k}={v}" for k, v in dataset_config["kwargs"].items())
|
|
print(f" - {dataset_config['name']}({kwargs_str})...", end=" ")
|
|
dataset_func(**dataset_config["kwargs"])
|
|
print("✓")
|
|
except Exception as e:
|
|
print(f"✗ (Error: {e})")
|
|
|
|
print("\nAll datasets processed!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
download_all_datasets()
|