Update app.py

app.py

@@ -1,16 +1,16 @@
 import streamlit as st
+import os
 import requests
 import pandas as pd
 import json
+import joblib  # Import joblib to load the local model
 
 # 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 = 'Add_Model_URL' 
-    token = 'Add_Databricks_Token' 
+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)
@@ -19,6 +19,18 @@ def score_model(dataset):
         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')
 
@@ -37,7 +49,7 @@ with st.form(key='churn_form'):
     salaryVec_0, salaryVec_1 = 0.0, 0.0
     if salary == 'Low':
         salaryVec_0 = 1.0
-    elif salary == 'Medium':
+    elif salary is 'Medium':
         salaryVec_1 = 1.0
 
     # Submit button
@@ -58,10 +70,20 @@ if submit_button:
         'salaryVec_1': [salaryVec_1]
     })
 
-    # Get prediction from the model
+    # Check if URL and token are specified
+    url = 'Add_Model_URL'
+    token = 'Add_Databricks_Token'
+    
     try:
-        prediction = score_model(input_data)
-        churn_prediction = prediction['predictions'][0]
+        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: