fairtfm 1.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- fairtfm-1.0.0/LICENSE +9 -0
- fairtfm-1.0.0/PKG-INFO +158 -0
- fairtfm-1.0.0/README.md +136 -0
- fairtfm-1.0.0/fairtfm/__init__.py +17 -0
- fairtfm-1.0.0/fairtfm/datasets/__init__.py +20 -0
- fairtfm-1.0.0/fairtfm/datasets/loaders.py +212 -0
- fairtfm-1.0.0/fairtfm/inference.py +153 -0
- fairtfm-1.0.0/fairtfm/model.py +360 -0
- fairtfm-1.0.0/fairtfm/utils.py +238 -0
- fairtfm-1.0.0/fairtfm.egg-info/PKG-INFO +158 -0
- fairtfm-1.0.0/fairtfm.egg-info/SOURCES.txt +14 -0
- fairtfm-1.0.0/fairtfm.egg-info/dependency_links.txt +1 -0
- fairtfm-1.0.0/fairtfm.egg-info/requires.txt +12 -0
- fairtfm-1.0.0/fairtfm.egg-info/top_level.txt +1 -0
- fairtfm-1.0.0/pyproject.toml +33 -0
- fairtfm-1.0.0/setup.cfg +4 -0
fairtfm-1.0.0/LICENSE
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
|
|
6
|
+
|
|
7
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
8
|
+
|
|
9
|
+
THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
fairtfm-1.0.0/PKG-INFO
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: fairtfm
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: A tabular foundation model pretrained with fairness constraints. It can make fair predictions in one forward pass.
|
|
5
|
+
Author-email: patrikken <kenfackjoslin@gmail.com>
|
|
6
|
+
Project-URL: Homepage, https://huggingface.co/patrikken/FairTFM
|
|
7
|
+
Requires-Python: >=3.9
|
|
8
|
+
Description-Content-Type: text/markdown
|
|
9
|
+
License-File: LICENSE
|
|
10
|
+
Requires-Dist: torch>=2.9
|
|
11
|
+
Requires-Dist: numpy>=1.24.0
|
|
12
|
+
Requires-Dist: pandas>=2.0.0
|
|
13
|
+
Requires-Dist: scikit-learn>=1.3.0
|
|
14
|
+
Requires-Dist: fairlearn>=0.13.0
|
|
15
|
+
Requires-Dist: huggingface_hub>=0.20
|
|
16
|
+
Provides-Extra: examples
|
|
17
|
+
Requires-Dist: xgboost>=3.2.0; extra == "examples"
|
|
18
|
+
Requires-Dist: seaborn>=0.13.2; extra == "examples"
|
|
19
|
+
Requires-Dist: folktables>=0.0.12; extra == "examples"
|
|
20
|
+
Requires-Dist: gdown; extra == "examples"
|
|
21
|
+
Dynamic: license-file
|
|
22
|
+
|
|
23
|
+
# FairTFM - Pretrained Tabular Foundation Model for fair predictions on Tabular Data
|
|
24
|
+
|
|
25
|
+
This repository provides **inference-only** FairTFM. The training code will follow.
|
|
26
|
+
|
|
27
|
+
## Quick Start
|
|
28
|
+
|
|
29
|
+
Install directly from source:
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
pip install .
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
or, for an editable/development install:
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
pip install -e .
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
This installs the `fairtfm` package along with its runtime dependencies (torch, numpy, pandas, scikit-learn, fairlearn, folktables, huggingface_hub). Extra dependencies used only by the paper reproduction / example scripts (xgboost, seaborn, gdown) can be installed with:
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
pip install -e ".[examples]"
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
### Model checkpoints
|
|
48
|
+
|
|
49
|
+
By default, `FairTFMClassifier()` automatically downloads the default checkpoint (λ = 0.7) from the [Hugging Face Hub](https://huggingface.co/patrikken/FairTFM) — no manual download needed:
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from fairtfm import FairTFMClassifier
|
|
53
|
+
|
|
54
|
+
classifier = FairTFMClassifier() # downloads FairTFM-0.7-epoch_10000.pt from Hugging Face
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
You can also point it at a specific checkpoint, either a local file, or a full Hugging Face URL:
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
classifier = FairTFMClassifier(model="https://huggingface.co/patrikken/FairTFM/blob/main/FairTFM-25-epoch_10000.pt")
|
|
61
|
+
classifier = FairTFMClassifier(model="path/to/local/checkpoint.pt")
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
All checkpoints produced during training across lambda values (0.7, 1.0, 10, 25), which are used to generate the fairness/accuracy Pareto front, are also bundled and available for downloading from [this](https://drive.google.com/uc?export=download&id=1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7) google drive link. Higher λ trades predictive performance for lower fairness-metric disparity, so depending on your use case a different checkpoint may give a stronger fairness/accuracy tradeoff — download the bundle and select the checkpoint that fits your needs:
|
|
65
|
+
|
|
66
|
+
```sh
|
|
67
|
+
pip install gdown
|
|
68
|
+
gdown 1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7
|
|
69
|
+
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
or use curl
|
|
73
|
+
|
|
74
|
+
```
|
|
75
|
+
wget --no-check-certificate "https://drive.google.com/uc?export=download&id=1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7" -o checkpoints.zip
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
### FairTFMClassifier interface overview
|
|
79
|
+
|
|
80
|
+
```python
|
|
81
|
+
from fairtfm import FairTFMClassifier, compute_fairness_metrics
|
|
82
|
+
|
|
83
|
+
# Load the default checkpoint from the Hugging Face Hub
|
|
84
|
+
classifier = FairTFMClassifier()
|
|
85
|
+
|
|
86
|
+
# or load a specific checkpoint (local path, Hugging Face repo id/URL)
|
|
87
|
+
classifier = FairTFMClassifier(model="path/to/checkpoint")
|
|
88
|
+
|
|
89
|
+
# Fit on training data
|
|
90
|
+
classifier.fit(X_train, y_train, s_train)
|
|
91
|
+
|
|
92
|
+
# Predict
|
|
93
|
+
predictions = classifier.predict(X_test, s_test)
|
|
94
|
+
probabilities = classifier.predict_proba(X_test, s_test)
|
|
95
|
+
|
|
96
|
+
# Get embeddings of the testing data
|
|
97
|
+
embeddings = classifier.transform(X_test, s_test)
|
|
98
|
+
|
|
99
|
+
# Fairness metrics (returns dict with performance metrics)
|
|
100
|
+
compute_fairness_metrics(X_test, y_test, s_test)
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
### ACSPumsDataset
|
|
104
|
+
|
|
105
|
+
```python
|
|
106
|
+
from fairtfm.datasets import ACSPumsDataset
|
|
107
|
+
|
|
108
|
+
dataset = ACSPumsDataset(
|
|
109
|
+
acs_task="acs_income", # Income, employment, mobility, travel_time, public_coverage
|
|
110
|
+
states=["CA"], # State codes
|
|
111
|
+
sensitive_attr_name="SEX" # SEX, RAC1P (Race), AGEP (Age)
|
|
112
|
+
)
|
|
113
|
+
dataset.preprocess()
|
|
114
|
+
X_train, X_test, y_train, y_test, s_train, s_test = dataset.get_splits()
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
## Supported Tasks
|
|
119
|
+
|
|
120
|
+
- `acs_income` - Income prediction
|
|
121
|
+
- `acs_employment` - Employment status
|
|
122
|
+
- `acs_mobility` - Geographic mobility
|
|
123
|
+
- `acs_public_coverage` - Public health insurance
|
|
124
|
+
- `acs_travel_time` - Travel time to work
|
|
125
|
+
|
|
126
|
+
## Sensitive Attributes
|
|
127
|
+
|
|
128
|
+
- `SEX` - Gender
|
|
129
|
+
- `RAC1P` - Race (White/Black)
|
|
130
|
+
- `AGEP` - Age (binarized by median)
|
|
131
|
+
|
|
132
|
+
## Code example
|
|
133
|
+
For inference example use the [notebook](notebook.ipynb) or the [inference_example.py](inference_example.py)
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
# Reproducing the paper main results.
|
|
137
|
+
|
|
138
|
+
For reproducing the paper's results run:
|
|
139
|
+
|
|
140
|
+
```python
|
|
141
|
+
python paper_results.py --full-eval
|
|
142
|
+
```
|
|
143
|
+
|
|
144
|
+
This will evaluate all the models in [`eval_config/eval_models.csv`](eval_config/eval_models.csv) on all 120 tasks in [`eval_config/fairness_tasks_eval.csv`](eval_config/fairness_tasks_eval.csv).
|
|
145
|
+
|
|
146
|
+
## Citation
|
|
147
|
+
If you use FairTFM for research purposes, please cite our [paper](https://openreview.net/forum?id=ajIvCEbadL):
|
|
148
|
+
|
|
149
|
+
```bibtex
|
|
150
|
+
@inproceedings{
|
|
151
|
+
kenfack2026training,
|
|
152
|
+
title={Training Fair Tabular Foundation Models},
|
|
153
|
+
author={Patrik Kenfack and Jesse C. Cresswell and Anthony L. Caterini and Samira Ebrahimi Kahou and Ulrich A{\"\i}vodji},
|
|
154
|
+
booktitle={2nd ICML Workshop on Foundation Models for Structured Data},
|
|
155
|
+
year={2026},
|
|
156
|
+
url={https://openreview.net/forum?id=ajIvCEbadL}
|
|
157
|
+
}
|
|
158
|
+
```
|
fairtfm-1.0.0/README.md
ADDED
|
@@ -0,0 +1,136 @@
|
|
|
1
|
+
# FairTFM - Pretrained Tabular Foundation Model for fair predictions on Tabular Data
|
|
2
|
+
|
|
3
|
+
This repository provides **inference-only** FairTFM. The training code will follow.
|
|
4
|
+
|
|
5
|
+
## Quick Start
|
|
6
|
+
|
|
7
|
+
Install directly from source:
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install .
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
or, for an editable/development install:
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install -e .
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
This installs the `fairtfm` package along with its runtime dependencies (torch, numpy, pandas, scikit-learn, fairlearn, folktables, huggingface_hub). Extra dependencies used only by the paper reproduction / example scripts (xgboost, seaborn, gdown) can be installed with:
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
pip install -e ".[examples]"
|
|
23
|
+
```
|
|
24
|
+
|
|
25
|
+
### Model checkpoints
|
|
26
|
+
|
|
27
|
+
By default, `FairTFMClassifier()` automatically downloads the default checkpoint (λ = 0.7) from the [Hugging Face Hub](https://huggingface.co/patrikken/FairTFM) — no manual download needed:
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
from fairtfm import FairTFMClassifier
|
|
31
|
+
|
|
32
|
+
classifier = FairTFMClassifier() # downloads FairTFM-0.7-epoch_10000.pt from Hugging Face
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
You can also point it at a specific checkpoint, either a local file, or a full Hugging Face URL:
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
classifier = FairTFMClassifier(model="https://huggingface.co/patrikken/FairTFM/blob/main/FairTFM-25-epoch_10000.pt")
|
|
39
|
+
classifier = FairTFMClassifier(model="path/to/local/checkpoint.pt")
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
All checkpoints produced during training across lambda values (0.7, 1.0, 10, 25), which are used to generate the fairness/accuracy Pareto front, are also bundled and available for downloading from [this](https://drive.google.com/uc?export=download&id=1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7) google drive link. Higher λ trades predictive performance for lower fairness-metric disparity, so depending on your use case a different checkpoint may give a stronger fairness/accuracy tradeoff — download the bundle and select the checkpoint that fits your needs:
|
|
43
|
+
|
|
44
|
+
```sh
|
|
45
|
+
pip install gdown
|
|
46
|
+
gdown 1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7
|
|
47
|
+
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
or use curl
|
|
51
|
+
|
|
52
|
+
```
|
|
53
|
+
wget --no-check-certificate "https://drive.google.com/uc?export=download&id=1SBztiK9SZZ_6-3I8oT8KadOR0JmuN-j7" -o checkpoints.zip
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
### FairTFMClassifier interface overview
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
from fairtfm import FairTFMClassifier, compute_fairness_metrics
|
|
60
|
+
|
|
61
|
+
# Load the default checkpoint from the Hugging Face Hub
|
|
62
|
+
classifier = FairTFMClassifier()
|
|
63
|
+
|
|
64
|
+
# or load a specific checkpoint (local path, Hugging Face repo id/URL)
|
|
65
|
+
classifier = FairTFMClassifier(model="path/to/checkpoint")
|
|
66
|
+
|
|
67
|
+
# Fit on training data
|
|
68
|
+
classifier.fit(X_train, y_train, s_train)
|
|
69
|
+
|
|
70
|
+
# Predict
|
|
71
|
+
predictions = classifier.predict(X_test, s_test)
|
|
72
|
+
probabilities = classifier.predict_proba(X_test, s_test)
|
|
73
|
+
|
|
74
|
+
# Get embeddings of the testing data
|
|
75
|
+
embeddings = classifier.transform(X_test, s_test)
|
|
76
|
+
|
|
77
|
+
# Fairness metrics (returns dict with performance metrics)
|
|
78
|
+
compute_fairness_metrics(X_test, y_test, s_test)
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
### ACSPumsDataset
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
from fairtfm.datasets import ACSPumsDataset
|
|
85
|
+
|
|
86
|
+
dataset = ACSPumsDataset(
|
|
87
|
+
acs_task="acs_income", # Income, employment, mobility, travel_time, public_coverage
|
|
88
|
+
states=["CA"], # State codes
|
|
89
|
+
sensitive_attr_name="SEX" # SEX, RAC1P (Race), AGEP (Age)
|
|
90
|
+
)
|
|
91
|
+
dataset.preprocess()
|
|
92
|
+
X_train, X_test, y_train, y_test, s_train, s_test = dataset.get_splits()
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
## Supported Tasks
|
|
97
|
+
|
|
98
|
+
- `acs_income` - Income prediction
|
|
99
|
+
- `acs_employment` - Employment status
|
|
100
|
+
- `acs_mobility` - Geographic mobility
|
|
101
|
+
- `acs_public_coverage` - Public health insurance
|
|
102
|
+
- `acs_travel_time` - Travel time to work
|
|
103
|
+
|
|
104
|
+
## Sensitive Attributes
|
|
105
|
+
|
|
106
|
+
- `SEX` - Gender
|
|
107
|
+
- `RAC1P` - Race (White/Black)
|
|
108
|
+
- `AGEP` - Age (binarized by median)
|
|
109
|
+
|
|
110
|
+
## Code example
|
|
111
|
+
For inference example use the [notebook](notebook.ipynb) or the [inference_example.py](inference_example.py)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
# Reproducing the paper main results.
|
|
115
|
+
|
|
116
|
+
For reproducing the paper's results run:
|
|
117
|
+
|
|
118
|
+
```python
|
|
119
|
+
python paper_results.py --full-eval
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
This will evaluate all the models in [`eval_config/eval_models.csv`](eval_config/eval_models.csv) on all 120 tasks in [`eval_config/fairness_tasks_eval.csv`](eval_config/fairness_tasks_eval.csv).
|
|
123
|
+
|
|
124
|
+
## Citation
|
|
125
|
+
If you use FairTFM for research purposes, please cite our [paper](https://openreview.net/forum?id=ajIvCEbadL):
|
|
126
|
+
|
|
127
|
+
```bibtex
|
|
128
|
+
@inproceedings{
|
|
129
|
+
kenfack2026training,
|
|
130
|
+
title={Training Fair Tabular Foundation Models},
|
|
131
|
+
author={Patrik Kenfack and Jesse C. Cresswell and Anthony L. Caterini and Samira Ebrahimi Kahou and Ulrich A{\"\i}vodji},
|
|
132
|
+
booktitle={2nd ICML Workshop on Foundation Models for Structured Data},
|
|
133
|
+
year={2026},
|
|
134
|
+
url={https://openreview.net/forum?id=ajIvCEbadL}
|
|
135
|
+
}
|
|
136
|
+
```
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""
|
|
2
|
+
FairTFM - Fair Tabular Foundation Model
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
from .inference import FairTFMClassifier
|
|
6
|
+
|
|
7
|
+
__version__ = "1.0.0"
|
|
8
|
+
__all__ = [
|
|
9
|
+
"FairTFMClassifier",
|
|
10
|
+
]
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def __getattr__(name):
|
|
14
|
+
if name == "ACSPumsDataset":
|
|
15
|
+
from .datasets import ACSPumsDataset
|
|
16
|
+
return ACSPumsDataset
|
|
17
|
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Fairness dataset loaders for ACS PUMS and other benchmark tasks.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
__all__ = [
|
|
6
|
+
"ACSPumsDataset",
|
|
7
|
+
]
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def __getattr__(name):
|
|
11
|
+
if name == "ACSPumsDataset":
|
|
12
|
+
try:
|
|
13
|
+
from .loaders import ACSPumsDataset
|
|
14
|
+
except ImportError as e:
|
|
15
|
+
raise ImportError(
|
|
16
|
+
"ACSPumsDataset requires the 'folktables' package. "
|
|
17
|
+
"Install it with `pip install fairtfm[examples]` or `pip install folktables`."
|
|
18
|
+
) from e
|
|
19
|
+
return ACSPumsDataset
|
|
20
|
+
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
|
@@ -0,0 +1,212 @@
|
|
|
1
|
+
"""
|
|
2
|
+
ACS PUMS dataset loader for fairness tasks.
|
|
3
|
+
Used for testing inference on real fairness benchmark datasets.
|
|
4
|
+
"""
|
|
5
|
+
import pandas as pd
|
|
6
|
+
import numpy as np
|
|
7
|
+
from sklearn.model_selection import train_test_split
|
|
8
|
+
from sklearn.preprocessing import StandardScaler, LabelEncoder
|
|
9
|
+
from folktables import (
|
|
10
|
+
ACSDataSource,
|
|
11
|
+
ACSEmployment,
|
|
12
|
+
ACSIncome,
|
|
13
|
+
ACSMobility,
|
|
14
|
+
ACSPublicCoverage,
|
|
15
|
+
ACSTravelTime,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class ACSPumsDataset:
|
|
20
|
+
"""
|
|
21
|
+
Loads American Community Survey (ACS) Public Use Microdata Sample (PUMS) tasks.
|
|
22
|
+
Provides fairness benchmarks with configurable sensitive attributes.
|
|
23
|
+
|
|
24
|
+
Example:
|
|
25
|
+
>>> dataset = ACSPumsDataset(
|
|
26
|
+
... acs_task="acs_income",
|
|
27
|
+
... states=["CA"],
|
|
28
|
+
... sensitive_attr_name="SEX"
|
|
29
|
+
... )
|
|
30
|
+
>>> dataset.preprocess()
|
|
31
|
+
>>> X_train, X_test, y_train, y_test, s_train, s_test = dataset.get_splits()
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
def __init__(
|
|
35
|
+
self,
|
|
36
|
+
acs_task: str = "acs_income",
|
|
37
|
+
states: list = None,
|
|
38
|
+
survey_year: int = 2018,
|
|
39
|
+
sensitive_attr_name: str = "SEX",
|
|
40
|
+
horizon: str = "1-Year",
|
|
41
|
+
seed_everything: int = 42,
|
|
42
|
+
max_samples: int = None,
|
|
43
|
+
test_size: float = 0.2
|
|
44
|
+
):
|
|
45
|
+
"""
|
|
46
|
+
Initialize ACS PUMS dataset loader.
|
|
47
|
+
|
|
48
|
+
Args:
|
|
49
|
+
acs_task: Task name ('acs_income', 'acs_employment', 'acs_mobility',
|
|
50
|
+
'acs_public_coverage', 'acs_travel_time')
|
|
51
|
+
states: List of state codes to include (e.g., ['CA', 'NY'])
|
|
52
|
+
survey_year: Survey year (default 2018)
|
|
53
|
+
sensitive_attr_name: Sensitive attribute ('SEX', 'RAC1P', 'AGEP')
|
|
54
|
+
horizon: Survey horizon ('1-Year' or '5-Year')
|
|
55
|
+
seed_everything: Random seed for reproducibility
|
|
56
|
+
max_samples: Maximum samples to load (None = all)
|
|
57
|
+
test_size: Proportion of data to use for testing
|
|
58
|
+
"""
|
|
59
|
+
self.sensitive_attr_name = sensitive_attr_name
|
|
60
|
+
self.acs_task = acs_task
|
|
61
|
+
self.states = states if states is not None else ["AL"]
|
|
62
|
+
self.survey_year = survey_year
|
|
63
|
+
self.horizon = horizon
|
|
64
|
+
self.seed = seed_everything
|
|
65
|
+
self.max_samples = max_samples
|
|
66
|
+
self.test_size = test_size
|
|
67
|
+
|
|
68
|
+
# Validation
|
|
69
|
+
if sensitive_attr_name not in ["SEX", "RAC1P", "AGEP"]:
|
|
70
|
+
raise ValueError(f"Unsupported sensitive attribute: {sensitive_attr_name}")
|
|
71
|
+
|
|
72
|
+
if acs_task not in [
|
|
73
|
+
"acs_income", "acs_employment", "acs_mobility",
|
|
74
|
+
"acs_public_coverage", "acs_travel_time",
|
|
75
|
+
]:
|
|
76
|
+
raise ValueError(f"Unsupported ACS task: {acs_task}")
|
|
77
|
+
|
|
78
|
+
state_name = "_".join(self.states) if self.states else "all_states"
|
|
79
|
+
self.name = f"{acs_task}_{state_name}_{sensitive_attr_name}"
|
|
80
|
+
|
|
81
|
+
self.data = None
|
|
82
|
+
self.target = None
|
|
83
|
+
self.sensitive_attr = None
|
|
84
|
+
self.train_idx = None
|
|
85
|
+
self.test_idx = None
|
|
86
|
+
self.X_train = None
|
|
87
|
+
self.X_test = None
|
|
88
|
+
self.y_train = None
|
|
89
|
+
self.y_test = None
|
|
90
|
+
self.s_train = None
|
|
91
|
+
self.s_test = None
|
|
92
|
+
|
|
93
|
+
def _get_task_class(self):
|
|
94
|
+
"""Get the appropriate task class for the ACS task."""
|
|
95
|
+
task_map = {
|
|
96
|
+
"acs_income": ACSIncome,
|
|
97
|
+
"acs_employment": ACSEmployment,
|
|
98
|
+
"acs_mobility": ACSMobility,
|
|
99
|
+
"acs_public_coverage": ACSPublicCoverage,
|
|
100
|
+
"acs_travel_time": ACSTravelTime,
|
|
101
|
+
}
|
|
102
|
+
return task_map[self.acs_task]
|
|
103
|
+
|
|
104
|
+
def preprocess(self):
|
|
105
|
+
"""
|
|
106
|
+
Download and preprocess the ACS PUMS dataset.
|
|
107
|
+
Applies task-specific transformations and handles sensitive attributes.
|
|
108
|
+
"""
|
|
109
|
+
# Download data
|
|
110
|
+
data_source = ACSDataSource(
|
|
111
|
+
survey_year=self.survey_year,
|
|
112
|
+
horizon=self.horizon,
|
|
113
|
+
survey="person",
|
|
114
|
+
root_dir="data"
|
|
115
|
+
)
|
|
116
|
+
acs_data = data_source.get_data(
|
|
117
|
+
states=self.states,
|
|
118
|
+
download=True,
|
|
119
|
+
random_seed=self.seed
|
|
120
|
+
)
|
|
121
|
+
|
|
122
|
+
# Subsample if needed
|
|
123
|
+
nb_sample = min(len(acs_data), self.max_samples) if self.max_samples else len(acs_data)
|
|
124
|
+
if self.max_samples is not None or len(self.states) == 0:
|
|
125
|
+
acs_data = acs_data.sample(n=nb_sample, random_state=self.seed)
|
|
126
|
+
|
|
127
|
+
# Apply task-specific transformations
|
|
128
|
+
task_class = self._get_task_class()
|
|
129
|
+
acs_data, self.target, _ = task_class.df_to_pandas(acs_data)
|
|
130
|
+
|
|
131
|
+
# Extract sensitive attributes
|
|
132
|
+
self.sensitive_attr = acs_data[self.sensitive_attr_name]
|
|
133
|
+
|
|
134
|
+
# Keep only relevant features
|
|
135
|
+
self.data = acs_data.drop(columns=[self.sensitive_attr_name])
|
|
136
|
+
|
|
137
|
+
# Handle race attribute
|
|
138
|
+
if self.sensitive_attr_name == "RAC1P":
|
|
139
|
+
idx = acs_data[self.sensitive_attr_name].isin([1, 2]) # White and Black only
|
|
140
|
+
self.data = self.data[idx]
|
|
141
|
+
self.target = self.target[idx]
|
|
142
|
+
self.sensitive_attr = self.sensitive_attr[idx]
|
|
143
|
+
|
|
144
|
+
# binary age attribute
|
|
145
|
+
elif self.sensitive_attr_name == "AGEP":
|
|
146
|
+
self.sensitive_attr = self.sensitive_attr.apply(
|
|
147
|
+
lambda x: 1 if x > 25 else 0
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
# Create train/test split
|
|
151
|
+
self.train_idx, self.test_idx = train_test_split(
|
|
152
|
+
range(len(self.data)),
|
|
153
|
+
test_size=self.test_size,
|
|
154
|
+
random_state=self.seed,
|
|
155
|
+
stratify=self.target,
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
self.data = pd.get_dummies(self.data, drop_first=True).values.astype(np.float32)
|
|
159
|
+
y_sc = LabelEncoder()
|
|
160
|
+
self.target = y_sc.fit_transform(self.target.squeeze())
|
|
161
|
+
|
|
162
|
+
s_sc = LabelEncoder()
|
|
163
|
+
self.sensitive_attr = s_sc.fit_transform(self.sensitive_attr.squeeze())
|
|
164
|
+
|
|
165
|
+
self.train_idx, self.test_idx = train_test_split(
|
|
166
|
+
range(len(self.data)),
|
|
167
|
+
test_size=self.test_size,
|
|
168
|
+
random_state=self.seed,
|
|
169
|
+
stratify=self.target,
|
|
170
|
+
)
|
|
171
|
+
self.input_dim = self.data.shape[1]
|
|
172
|
+
return self
|
|
173
|
+
|
|
174
|
+
def get_splits(self):
|
|
175
|
+
"""
|
|
176
|
+
Get train/test splits with features, labels, and sensitive attributes.
|
|
177
|
+
|
|
178
|
+
Returns:
|
|
179
|
+
Tuple of (X_train, X_test, y_train, y_test, s_train, s_test)
|
|
180
|
+
"""
|
|
181
|
+
if self.data is None or self.target is None or self.sensitive_attr is None:
|
|
182
|
+
raise ValueError("Dataset not preprocessed. Call preprocess() first.")
|
|
183
|
+
scaler = StandardScaler()
|
|
184
|
+
""" self.data = pd.DataFrame(
|
|
185
|
+
scaler.fit_transform(self.data), columns=self.data.columns
|
|
186
|
+
) """
|
|
187
|
+
|
|
188
|
+
scaler = StandardScaler()
|
|
189
|
+
|
|
190
|
+
self.X_train = self.data[self.train_idx]
|
|
191
|
+
self.X_test = self.data[self.test_idx]
|
|
192
|
+
self.y_train = self.target[self.train_idx]
|
|
193
|
+
self.y_test = self.target[self.test_idx]
|
|
194
|
+
self.s_train = self.sensitive_attr[self.train_idx]
|
|
195
|
+
self.s_test = self.sensitive_attr[self.test_idx]
|
|
196
|
+
|
|
197
|
+
self.X_train = scaler.fit_transform(self.X_train).astype("float32")
|
|
198
|
+
self.X_test = scaler.transform(self.X_test).astype("float32")
|
|
199
|
+
return (
|
|
200
|
+
self.X_train, self.X_test,
|
|
201
|
+
self.y_train, self.y_test,
|
|
202
|
+
self.s_train, self.s_test
|
|
203
|
+
)
|
|
204
|
+
|
|
205
|
+
def __len__(self):
|
|
206
|
+
"""Return total number of samples."""
|
|
207
|
+
if self.data is None:
|
|
208
|
+
return 0
|
|
209
|
+
return len(self.data)
|
|
210
|
+
|
|
211
|
+
def __repr__(self):
|
|
212
|
+
return f"ACSPumsDataset({self.name}, n_samples={len(self)})"
|