Update utils/inference.py

utils/inference.py

@@ -1,22 +1,30 @@
 import numpy as np
 import pandas as pd
 import joblib
+import streamlit as st
 from catboost import CatBoostClassifier
 from xgboost import XGBClassifier
 
 MODEL_DIR = "models"
+THRESHOLD_S1 = 0.463
 
-model_s1 = CatBoostClassifier()
-model_s1.load_model(f"{MODEL_DIR}/stage1_catboost.cbm")
 
-model_s2 = XGBClassifier()
-model_s2.load_model(f"{MODEL_DIR}/stage2_xgb.json")
+@st.cache_resource
+def load_models():
+    model_s1 = CatBoostClassifier()
+    model_s1.load_model(f"{MODEL_DIR}/stage1_catboost.cbm")
 
-encoder = joblib.load(f"{MODEL_DIR}/encoder_s2.pkl")
+    model_s2 = XGBClassifier()
+    model_s2.load_model(f"{MODEL_DIR}/stage2_xgb.json")
+
+    encoder = joblib.load(f"{MODEL_DIR}/encoder_s2.pkl")
+
+    return model_s1, model_s2, encoder
 
-THRESHOLD_S1 = 0.463
 
 def run_inference(X_emb, meta, threshold):
+    model_s1, model_s2, encoder = load_models()
+
     df_meta = pd.DataFrame(meta)
 
     probs_s1 = model_s1.predict_proba(X_emb)[:, 1]
@@ -31,4 +39,4 @@ def run_inference(X_emb, meta, threshold):
     df_pass["Final_Score"] = final_score
     df_pass["Final_Prediction"] = (final_score >= threshold).astype(int)
 
-    return df_pass
+    return df_pass