Spaces:
Running
Running
import streamlit as st | |
import pandas as pd | |
import lightgbm as lgb | |
import pickle | |
# Load the trained LightGBM model | |
with open('lgb.pkl', 'rb') as model_file: | |
model = pickle.load(model_file) | |
# Define mappings | |
workclass_map = { | |
'Private': 1, | |
'State-gov': 2, | |
'Federal-gov': 3, | |
'Self-emp-not-inc': 4, | |
'Self-emp-inc': 5, | |
'Local-gov': 6, | |
'Without-pay': 7, | |
'Never-worked': 8, | |
'?': 9 | |
} | |
education_ranks = { | |
'Preschool': 1, | |
'1st-4th': 2, | |
'5th-6th': 3, | |
'7th-8th': 4, | |
'9th': 5, | |
'10th': 6, | |
'11th': 7, | |
'12th': 8, | |
'HS-grad': 9, | |
'Some-college': 10, | |
'Assoc-voc': 11, | |
'Assoc-acdm': 12, | |
'Bachelors': 13, | |
'Masters': 14, | |
'Prof-school': 15, | |
'Doctorate': 16 | |
} | |
marital_map = { | |
'Married-civ-spouse': 1, | |
'Married-spouse-absent': 1, | |
'Married-AF-spouse': 1, | |
'Widowed': 2, | |
'Divorced': 2, | |
'Separated': 2, | |
'Never-married': 2 | |
} | |
occupation_map = { | |
'Exec-managerial': 1, | |
'Machine-op-inspct': 2, | |
'Prof-specialty': 3, | |
'Other-service': 4, | |
'Adm-clerical': 5, | |
'Craft-repair': 6, | |
'Transport-moving': 7, | |
'Handlers-cleaners': 8, | |
'Sales': 9, | |
'Farming-fishing': 10, | |
'Tech-support': 11, | |
'Protective-serv': 12, | |
'Armed-Forces': 13, | |
'Priv-house-serv': 14 | |
} | |
relationship_map = { | |
'Not-in-family': 1, | |
'Unmarried': 2, | |
'Own-child': 3, | |
'Other-relative': 4, | |
'Husband': 5, | |
'Wife': 6 | |
} | |
income_map = { | |
'<=50K': 0, | |
'>50K': 1 | |
} | |
# Define the input fields for the user to provide data | |
def get_user_input(): | |
age = st.number_input('Age', min_value=0, max_value=120, value=30) | |
workclass = st.selectbox('Workclass', list(workclass_map.keys())) | |
fnlwgt = st.number_input('Fnlwgt', min_value=0, value=100000) | |
education = st.selectbox('Education', list(education_ranks.keys())) | |
education_num = st.number_input('Education Num', min_value=0, max_value=20, value=10) | |
marital_status = st.selectbox('Marital Status', list(marital_map.keys())) | |
occupation = st.selectbox('Occupation', list(occupation_map.keys())) | |
relationship = st.selectbox('Relationship', list(relationship_map.keys())) | |
capital_gain = st.number_input('Capital Gain', min_value=0, value=0) | |
capital_loss = st.number_input('Capital Loss', min_value=0, value=0) | |
hours_per_week = st.number_input('Hours Per Week', min_value=0, max_value=168, value=40) | |
user_data = { | |
'age': age, | |
'workclass_rank': workclass_map[workclass], | |
'fnlwgt': fnlwgt, | |
'education_rank': education_ranks[education], | |
'education.num': education_num, | |
'marital_status_binary': marital_map[marital_status], | |
'occupation_rank': occupation_map[occupation], | |
'relationship_rank': relationship_map[relationship], | |
'capital.gain': capital_gain, | |
'capital.loss': capital_loss, | |
'hours.per.week': hours_per_week | |
} | |
features = pd.DataFrame(user_data, index=[0]) | |
return features | |
# Main function to run the app | |
def main(): | |
st.title('Income Prediction App') | |
st.write('This app predicts whether a person makes over $50K a year based on their demographics and work information.') | |
user_input = get_user_input() | |
st.subheader('User Input:') | |
st.write(user_input) | |
prediction = model.predict(user_input) | |
prediction_proba = model.predict_proba(user_input)[:, 1] | |
st.subheader('Prediction:') | |
if prediction[0] == 1: | |
st.write('The model predicts: Income > $50K') | |
else: | |
st.write('The model predicts: Income <= $50K') | |
st.subheader('Prediction Probability:') | |
st.write(f'Probability of making over $50K: {prediction_proba[0]:.2f}') | |
if __name__ == '__main__': | |
main() | |