Spaces:
Sleeping
Sleeping
File size: 2,700 Bytes
8f559f1 | 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 63 64 65 66 67 68 69 70 71 72 73 74 | """Prepare a MemisisLabs dataset for TabDiff (github.com/MinkaiXu/TabDiff).
Writes the CSV + Info JSON into a TabDiff checkout's data/ layout so you can run
`process_dataset.py`, then train/sample. Column-type indices (num/cat/target) are
auto-detected from our dataset registry.
Usage (on the lab server, inside the TabDiff repo's conda env):
python scripts/tabdiff_prepare.py --dataset openml_45040 \
--name schizophrenia --tabdiff-root /path/to/TabDiff
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from pipeline import datasets # noqa: E402
from pipeline.metadata import is_categorical # noqa: E402
def main() -> None:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--dataset", default="openml_45040", help="MemisisLabs dataset id")
ap.add_argument("--name", default="schizophrenia", help="TabDiff dataset name")
ap.add_argument("--tabdiff-root", required=True, help="path to the cloned TabDiff repo")
args = ap.parse_args()
real, target, _protected, task = datasets.load(args.dataset)
cols = list(real.columns)
target_idx = cols.index(target)
num_idx, cat_idx = [], []
for i, c in enumerate(cols):
if i == target_idx:
continue
(cat_idx if is_categorical(real[c]) else num_idx).append(i)
root = Path(args.tabdiff_root)
(root / "data" / args.name).mkdir(parents=True, exist_ok=True)
(root / "data" / "Info").mkdir(parents=True, exist_ok=True)
csv_path = root / "data" / args.name / f"{args.name}.csv"
real.to_csv(csv_path, index=False)
info = {
"name": args.name,
"task_type": "binclass" if task == "binclass" else "regression",
"header": "infer",
"column_names": None,
"num_col_idx": num_idx,
"cat_col_idx": cat_idx,
"target_col_idx": [target_idx],
"file_type": "csv",
"data_path": f"data/{args.name}/{args.name}.csv",
"test_path": None,
}
info_path = root / "data" / "Info" / f"{args.name}.json"
info_path.write_text(json.dumps(info, indent=4))
print(f"wrote {csv_path} ({real.shape})")
print(f"wrote {info_path}")
print(f" num_col_idx={num_idx}\n cat_col_idx={cat_idx}\n target_col_idx=[{target_idx}] ({target})")
print("\nNext (in the TabDiff repo):")
print(f" python process_dataset.py --dataname {args.name}")
print(f" python main.py --dataname {args.name} --mode train --no_wandb")
print(f" python main.py --dataname {args.name} --mode test --report --no_wandb")
if __name__ == "__main__":
main()
|