Chethan4638/proactive-churn-predictor

1import pandas as pd2import mlflow3from fastapi import FastAPI4from pydantic import BaseModel5 6# Initialize the FastAPI app7app = FastAPI()8 9# --- Load Model and Data ---10mlflow.set_tracking_uri("http://127.0.0.1:5000")11RUN_ID = "671601d8c7674c188193c47eb56462e4"12MODEL_URI = f"runs:/{RUN_ID}/model"13DB_FILE = 'data/model_input.csv'14 15# --- Global Variables ---16model = None17df_sorted_results = pd.DataFrame() # This will hold our pre-calculated list18 19# --- Startup Event ---20@app.on_event("startup")21def load_and_predict():22    """23    This function runs ONCE when the uvicorn server starts.24    It loads the model and pre-calculates all churn probabilities.25    """26    global model, df_sorted_results27    28    # 1. Load Model29    try:30        model = mlflow.sklearn.load_model(MODEL_URI)31        print("Model loaded successfully.")32    except Exception as e:33        print(f"Error loading model: {e}")34        return35 36    # 2. Load Data37    try:38        df_full_data = pd.read_csv(DB_FILE)39        print("Data loaded successfully.")40    except FileNotFoundError:41        print(f"Error: {DB_FILE} not found.")42        return43 44    # 3. Pre-calculate Predictions (The "heavy lifting")45    print("Pre-calculating all user probabilities...")46    try:47        user_ids = df_full_data['user_id']48        X_to_predict = df_full_data.drop(['user_id', 'churn'], axis=1)49        50        churn_probabilities = model.predict_proba(X_to_predict)[:, 1]51 52        # Create the results DataFrame53        df_results = pd.DataFrame({54            'user_id': user_ids,55            'churn_probability': churn_probabilities56        })57        58        # Sort it *once* and save to our global variable59        df_sorted_results = df_results.sort_values(by='churn_probability', ascending=False)60        print("All probabilities calculated and sorted. API is ready.")61        62    except Exception as e:63        print(f"Error during pre-calculation: {e}")64 65 66# --- API Endpoints ---67@app.get("/")68def read_root():69    return {"status": "Proactive Churn API is running."}70 71 72@app.get("/get-retention-list/")73def get_retention_list(top_n: int = 500):  # <-- !! HERE IS THE FIX !!74    """75    This is now super fast! It just slices the pre-calculated DataFrame.76    """77    if df_sorted_results.empty:78        return {"error": "Server is still initializing or failed to load data. Check API logs."}79 80    # Just take the top N rows from the list we already made81    df_top_n = df_sorted_results.head(top_n)82    83    # Convert to a dictionary (JSON-friendly) and return it84    return df_top_n.to_dict(orient='records')