Update app.py

app.py

@@ -54,6 +54,7 @@ def cleaning_features(data):
     }
 
     gender_mapping = {"Male":1,"Female":0}
+    country_mapping = {"United-States":1,"Other":0}
     
     numeric_cols = ['age', 'educational-num', 'hours-per-week']
     columns_to_encode = ['race','marital-status','relationship']
@@ -62,13 +63,11 @@ def cleaning_features(data):
     data['workclass'] = le.fit_transform(data['workclass'])
     data['occupation'] = le.fit_transform(data['occupation'])
     data['gender'] = data['gender'].map(gender_mapping)
+    data['native-country'] = data['native-country'].map(country_mapping)
     data['educational-num'] = data['educational-num'].map(education_num_mapping)
     
     data[numeric_cols] = scaler.fit_transform(data[numeric_cols])
-    
-    # Binarize native country
-    data['native-country'] = data['native-country'].apply(lambda x: x == 'United-States')
-    data['native-country'] = data['native-country'].astype(int)
+
     #data = pca(data)
     return data