text: import pandas as pd
  import numpy as np
  import paddle.fluid as fluid
  from args import *
  
  
  def fun1(x):
      if x == ' 50000+.':
          return 1
      else:
          return 0
  
  
  def fun2(x):
      if x == ' Never married':
          return 1
      else:
          return 0
  
  
  def data_preparation(train_path, test_path, train_data_path, test_data_path):
      # The column names are from
      # https://www2.1010data.com/documentationcenter/prod/Tutorials/MachineLearningExamples/CensusIncomeDataSet.html
      column_names = [
          'age', 'class_worker', 'det_ind_code', 'det_occ_code', 'education',
          'wage_per_hour', 'hs_college', 'marital_stat', 'major_ind_code',
          'major_occ_code', 'race', 'hisp_origin', 'sex', 'union_member',
          'unemp_reason', 'full_or_part_emp', 'capital_gains', 'capital_losses',
          'stock_dividends', 'tax_filer_stat', 'region_prev_res',
          'state_prev_res', 'det_hh_fam_stat', 'det_hh_summ', 'instance_weight',
          'mig_chg_msa', 'mig_chg_reg', 'mig_move_reg', 'mig_same',
          'mig_prev_sunbelt', 'num_emp', 'fam_under_18', 'country_father',
          'country_mother', 'country_self', 'citizenship', 'own_or_self',
          'vet_question', 'vet_benefits', 'weeks_worked', 'year', 'income_50k'
      ]
  
      # Load the dataset in Pandas
      train_df = pd.read_csv(
          train_path,
          delimiter=',',
          header=None,
          index_col=None,
          names=column_names)
      other_df = pd.read_csv(
          test_path,
          delimiter=',',
          header=None,
          index_col=None,
          names=column_names)
  
      # First group of tasks according to the paper
      label_columns = ['income_50k', 'marital_stat']
  
      # One-hot encoding categorical columns
      categorical_columns = [
          'class_worker', 'det_ind_code', 'det_occ_code', 'education',
          'hs_college', 'major_ind_code', 'major_occ_code', 'race', 'hisp_origin',
          'sex', 'union_member', 'unemp_reason', 'full_or_part_emp',
          'tax_filer_stat', 'region_prev_res', 'state_prev_res',
          'det_hh_fam_stat', 'det_hh_summ', 'mig_chg_msa', 'mig_chg_reg',
          'mig_move_reg', 'mig_same', 'mig_prev_sunbelt', 'fam_under_18',
          'country_father', 'country_mother', 'country_self', 'citizenship',
          'vet_question'
      ]
      train_raw_labels = train_df[label_columns]
      other_raw_labels = other_df[label_columns]
      transformed_train = pd.get_dummies(train_df, columns=categorical_columns)
      transformed_other = pd.get_dummies(other_df, columns=categorical_columns)
  
      # Filling the missing column in the other set
      transformed_other[
          'det_hh_fam_stat_ Grandchild <18 ever marr not in subfamily'] = 0
      # get label
      transformed_train['income_50k'] = transformed_train['income_50k'].apply(
          lambda x: fun1(x))
      transformed_train['marital_stat'] = transformed_train['marital_stat'].apply(
          lambda x: fun2(x))
      transformed_other['income_50k'] = transformed_other['income_50k'].apply(
          lambda x: fun1(x))
      transformed_other['marital_stat'] = transformed_other['marital_stat'].apply(
          lambda x: fun2(x))
      # Split the other dataset into 1:1 validation to test according to the paper
      validation_indices = transformed_other.sample(
          frac=0.5, replace=False, random_state=1).index
      test_indices = list(set(transformed_other.index) - set(validation_indices))
      validation_data = transformed_other.iloc[validation_indices]
      test_data = transformed_other.iloc[test_indices]
  
      cols = transformed_train.columns.tolist()
      cols.insert(0, cols.pop(cols.index('income_50k')))
      cols.insert(0, cols.pop(cols.index('marital_stat')))
      transformed_train = transformed_train[cols]
      test_data = test_data[cols]
      validation_data = validation_data[cols]
  
      print(transformed_train.shape, transformed_other.shape,
            validation_data.shape, test_data.shape)
      transformed_train.to_csv(train_data_path + 'train_data.csv', index=False)
      test_data.to_csv(test_data_path + 'test_data.csv', index=False)
  
  
  args = data_preparation_args()
  data_preparation(args.train_path, args.test_path, args.train_data_path,args.test_data_path)
  