Update app.py

app.py

@@ -32,6 +32,8 @@ def predict(age, workclass, education, marital_status, occupation, relationship,
     return "Income >50K" if prediction == 1 else "Income <=50K"
 
 def cleaning_features(data):
+    with open('race_onehot_encoder.pkl', 'rb') as enc_file:
+        encoder = pickle.load(enc_file)
     
     with open('label_encoder_work.pkl', 'rb') as le_file:
         le_work = pickle.load(le_file)
@@ -64,7 +66,8 @@ def cleaning_features(data):
     country_mapping = {"United-States":1,"Other":0}
     
     numeric_cols = ['age', 'educational-num', 'hours-per-week']
-    columns_to_encode = ['race','marital-status','relationship']
+    # columns_to_encode = ['race','marital-status','relationship']
+    columns_to_encode = ['race']
     
     data['workclass'] = le_work.transform(data['workclass'])
     data['occupation'] = le_occ.transform(data['occupation'])
@@ -74,6 +77,13 @@ def cleaning_features(data):
     
     data[numeric_cols] = scaler.transform(data[numeric_cols])
 
+    for N in columns_to_encode:
+        race_encoded = encoder.transform(data[[N]])
+        race_encoded_cols = encoder.get_feature_names_out([N])
+        race_encoded_df = pd.DataFrame(race_encoded, columns=race_encoded_cols, index=data.index)
+        # Combine the encoded data with original dataframe
+        data = pd.concat([data.drop(N, axis=1), race_encoded_df], axis=1)
+
     data = pca(data)
     return data