# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

import pandas as pd

# Read in csv but add in column names
df = pd.read_csv('../datasets/adult.csv',
                 names=['age', 'workclass', 'fnlwgt', 'education', 'education-num', 'marital-status',
                        'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss',
                        'hours-per-week', 'native-country', 'income'])

# Convert binary labels to 0/1 representation
df = df.replace(' <=50K', 0)
df = df.replace(' >50K', 1)

# Write to a new file
df.to_csv('../datasets/adult_cleaned.csv', index=False)
