cestwc/census-income

import pandas as pd
from datasets import Dataset, DatasetDict, Features, Value, ClassLabel

# URLs
url1 = "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.data"
url2 = "https://archive.ics.uci.edu/ml/machine-learning-databases/adult/adult.test"

# Column names
columns = [
    "age", "workclass", "fnlwgt", "education", "education-num", "marital-status",
    "occupation", "relationship", "race", "sex", "capital-gain", "capital-loss",
    "hours-per-week", "native-country", "income"
]


# Load datasets
df_train = pd.read_csv(url1, names=columns, skipinitialspace=True)
df_test = pd.read_csv(url2, names=columns, skipinitialspace=True, skiprows=1)

# Convert continuous columns to float
continuous_columns = ["age", "fnlwgt", "education-num", "capital-gain", "capital-loss", "hours-per-week"]
for col in continuous_columns:
    df_train[col] = pd.to_numeric(df_train[col], errors='coerce')
    df_test[col] = pd.to_numeric(df_test[col], errors='coerce')

df_test['income'] = df_test['income'].str.rstrip('.') # This is somewhat critical.

# Define categorical columns
categorical_columns = [
    "workclass", "education", "marital-status", "occupation", "relationship",
    "race", "sex", "native-country", "income"
]

# Dictionary to store category mappings
category_mappings = {}

for col in categorical_columns:
    # Convert train column to category and extract categories
    df_train[col] = df_train[col].astype("category")
    category_mappings[col] = df_train[col].cat.categories.to_list()  # Store category order

    # Apply the same category mapping to test
    df_test[col] = pd.Categorical(df_test[col], categories=category_mappings[col])

    # Convert to integer codes
    df_train[col] = df_train[col].cat.codes
    df_test[col] = df_test[col].cat.codes

# Define Hugging Face dataset schema
hf_features = Features({
    "age": Value("int64"),
    "workclass": ClassLabel(names=category_mappings["workclass"]),
    "fnlwgt": Value("int64"),
    "education": ClassLabel(names=category_mappings["education"]),
    "education-num": Value("int64"),
    "marital-status": ClassLabel(names=category_mappings["marital-status"]),
    "occupation": ClassLabel(names=category_mappings["occupation"]),
    "relationship": ClassLabel(names=category_mappings["relationship"]),
    "race": ClassLabel(names=category_mappings["race"]),
    "sex": ClassLabel(names=category_mappings["sex"]),
    "capital-gain": Value("int64"),
    "capital-loss": Value("int64"),
    "hours-per-week": Value("int64"),
    "native-country": ClassLabel(names=category_mappings["native-country"]),
    "income": ClassLabel(names=category_mappings["income"])
})

# Convert pandas DataFrame to Hugging Face Dataset
hf_train = Dataset.from_pandas(df_train, features=hf_features)
hf_test = Dataset.from_pandas(df_test, features=hf_features)

# Create a dataset dictionary
hf_dataset = DatasetDict({
    "train": hf_train,
    "test": hf_test
})

# Print dataset structure
print(hf_dataset)

DatasetDict({
    train: Dataset({
        features: ['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income'],
        num_rows: 32561
    })
    test: Dataset({
        features: ['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country', 'income'],
        num_rows: 16281
    })
})