Abhishek0323's picture
Update app.py
057e8cb verified
import streamlit as st
import os
import requests
import pandas as pd
import json
import joblib
# Function to create the payload for the model
def create_tf_serving_json(data):
return {'inputs': {name: data[name].tolist() for name in data.keys()} if isinstance(data, dict) else data.tolist()}
# Function to send a request to the model endpoint
def score_model(dataset, url, token):
headers = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}
ds_dict = {'dataframe_split': dataset.to_dict(orient='split')} if isinstance(dataset, pd.DataFrame) else create_tf_serving_json(dataset)
data_json = json.dumps(ds_dict, allow_nan=True)
response = requests.request(method='POST', headers=headers, url=url, data=data_json)
if response.status_code != 200:
raise Exception(f'Request failed with status {response.status_code}, {response.text}')
return response.json()
# Load local model function
def load_local_model():
return joblib.load('Random Forest - Grid Search_model.pkl')
# Function to make predictions with local model
def predict_with_local_model(model, input_data):
input_data = input_data[['satisfaction_level', 'last_evaluation', 'number_project',
'average_montly_hours', 'time_spend_company', 'Work_accident',
'promotion_last_5years', 'salaryVec_0', 'salaryVec_1']]
prediction = model.predict(input_data)
return prediction[0]
# Streamlit app UI
st.title('Employee Churn Prediction')
# Create a form
with st.form(key='churn_form'):
satisfaction_level = st.slider('Satisfaction Level', 0.0, 1.0, 0.5)
last_evaluation = st.slider('Last Evaluation', 0.0, 1.0, 0.5)
number_project = st.slider('Number of Projects', 1, 10, 3)
average_montly_hours = st.slider('Average Monthly Hours', 50, 350, 200)
time_spend_company = st.slider('Time Spent in Company (years)', 1, 10, 3)
work_accident = st.selectbox('Work Accident', [0, 1])
promotion_last_5years = st.selectbox('Promotion in Last 5 Years', [0, 1])
salary = st.selectbox('Salary Level', ['Low', 'Medium', 'High'])
# Encode salary into one-hot vectors
salaryVec_0, salaryVec_1 = 0.0, 0.0
if salary == 'Low':
salaryVec_0 = 1.0
elif salary is 'Medium':
salaryVec_1 = 1.0
# Submit button
submit_button = st.form_submit_button(label='Predict Churn')
# Handle form submission
if submit_button:
# Create a DataFrame with the input data
input_data = pd.DataFrame({
'satisfaction_level': [satisfaction_level],
'last_evaluation': [last_evaluation],
'number_project': [number_project],
'average_montly_hours': [average_montly_hours],
'time_spend_company': [time_spend_company],
'Work_accident': [work_accident],
'promotion_last_5years': [promotion_last_5years],
'salaryVec_0': [salaryVec_0],
'salaryVec_1': [salaryVec_1]
})
# Check if URL and token are specified
url = 'Add_Model_URL'
token = 'Add_Databricks_Token'
try:
if url != 'Add_Model_URL' and token != 'Add_Databricks_Token':
# Use the API if URL and token are specified
prediction = score_model(input_data, url, token)
churn_prediction = prediction['predictions'][0]
else:
# Use the local model if URL and token are not specified
local_model = load_local_model()
churn_prediction = predict_with_local_model(local_model, input_data)
if churn_prediction == 1:
st.write('The employee is likely to churn.')
else:
st.write('The employee is not likely to churn.')
except Exception as e:
st.error(f'Error: {e}')