Initial commit

inference.py

@@ -0,0 +1,37 @@
+import joblib
+import pandas as pd
+from lifelines import CoxPHFitter
+
+class HRAttritionModel:
+    def __init__(self, model_path):
+        self.model = joblib.load(model_path)
+        self.features = ['Age', 'DistanceFromHome', 'Education', 'NumCompaniesWorked', 'PercentSalaryHike', 
+                         'TotalWorkingYears', 'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsInCurrentRole', 
+                         'YearsSinceLastPromotion', 'YearsWithCurrManager', 'BusinessTravel_Travel_Rarely', 
+                         'BusinessTravel_Travel_Frequently', 'Department_Research & Development', 'Department_Sales', 
+                         'EducationField_Life Sciences', 'EducationField_Medical', 'EducationField_Marketing', 
+                         'EducationField_Other', 'EducationField_Technical Degree', 'Gender_Male', 'JobRole_Research Scientist', 
+                         'JobRole_Sales Executive', 'JobRole_Laboratory Technician', 'JobRole_Manufacturing Director', 
+                         'JobRole_Healthcare Representative', 'JobRole_Manager', 'JobRole_Sales Representative', 
+                         'JobRole_Research Director', 'MaritalStatus_Married', 'MaritalStatus_Single', 'OverTime_Yes']
+    
+    def predict_survival(self, input_data):
+        df = pd.DataFrame([input_data], columns=self.features)
+        survival_function = self.model.predict_survival_function(df)
+        return survival_function.T
+
+# Load the model and make a prediction for testing
+if __name__ == "__main__":
+    model = HRAttritionModel('cox_model.pkl')
+    sample_input = {'Age': 41, 'DistanceFromHome': 1, 'Education': 2, 'NumCompaniesWorked': 1, 'PercentSalaryHike': 11, 
+                    'TotalWorkingYears': 8, 'TrainingTimesLastYear': 0, 'WorkLifeBalance': 1, 'YearsInCurrentRole': 4, 
+                    'YearsSinceLastPromotion': 0, 'YearsWithCurrManager': 5, 'BusinessTravel_Travel_Rarely': 1, 
+                    'BusinessTravel_Travel_Frequently': 0, 'Department_Research & Development': 0, 'Department_Sales': 1, 
+                    'EducationField_Life Sciences': 1, 'EducationField_Medical': 0, 'EducationField_Marketing': 0, 
+                    'EducationField_Other': 0, 'EducationField_Technical Degree': 0, 'Gender_Male': 1, 
+                    'JobRole_Research Scientist': 0, 'JobRole_Sales Executive': 0, 'JobRole_Laboratory Technician': 0, 
+                    'JobRole_Manufacturing Director': 0, 'JobRole_Healthcare Representative': 0, 'JobRole_Manager': 0, 
+                    'JobRole_Sales Representative': 0, 'JobRole_Research Director': 0, 'MaritalStatus_Married': 0, 
+                    'MaritalStatus_Single': 1, 'OverTime_Yes': 0}
+    prediction = model.predict_survival(sample_input)
+    print(prediction)

requirements.txt

@@ -0,0 +1,3 @@
+pandas
+joblib
+lifelines