fastapi_hf / routes /ML_LinearRegression_ChurnPredictor.py
looh2's picture
Remove unused embedding files and ingestion script; refactor PDF upload to handle embeddings directly
62d37ba
Raw
History Blame Contribute Delete
2.58 kB
from fastapi import APIRouter
from pydantic import BaseModel
import joblib
import pandas as pd
from .config_huggingface import build_model_url, download_artifact_if_needed
router = APIRouter(tags=["Machine Learning"])
# Define the request model for linear regression
class LinearRegressionRequest(BaseModel):
age: int = 30
monthly_spend: int = 50
tenure_months: int = 12
from typing import Optional, Any
MODEL_STATE: dict[str, Optional[Any]] = {
"model": None,
"error": None,
}
MODEL_URL = build_model_url("ML_LinearRegression_ChurnPredictor.joblib")
def _ensure_model_loaded() -> None:
if MODEL_STATE["model"] is not None:
return
try:
model_path = download_artifact_if_needed(MODEL_URL)
MODEL_STATE["model"] = joblib.load(model_path)
MODEL_STATE["error"] = None
except Exception as e:
MODEL_STATE["error"] = str(e)
raise
@router.post('/models/linear_regression', summary="Predict churn risk with Bayesian Ridge Regression")
def predict_linear_regression(data: LinearRegressionRequest):
import traceback
try:
_ensure_model_loaded()
except Exception:
detail = "Model not loaded."
if MODEL_STATE["error"]:
detail = f"Model not loaded: {MODEL_STATE['error']}"
# Log traceback for debugging
print("Model load error:", MODEL_STATE["error"])
print(traceback.format_exc())
return {"error": detail, "traceback": traceback.format_exc(), "status": 500}
model = MODEL_STATE["model"]
if model is None:
# Defensive: model still not loaded
error_msg = f"Model is None after loading. Error: {MODEL_STATE['error']}"
print(error_msg)
return {"error": error_msg, "status": 500}
# Convert input data to DataFrame matching training features
new_customer_data = pd.DataFrame([[data.age, data.monthly_spend, data.tenure_months]],
columns=['age', 'monthly_spend', 'tenure_months'])
# Make prediction
try:
pred_mean = model.predict(new_customer_data)[0]
except Exception as e:
print("Prediction error:", str(e))
print(traceback.format_exc())
return {"error": f"Prediction failed: {str(e)}", "traceback": traceback.format_exc(), "status": 500}
percent = round(float(pred_mean) * 100, 1)
if percent > 100:
percent_str = '>100%'
elif percent < 1:
percent_str = '<1%'
else:
percent_str = f"{percent}%"
return {
"predicted_churn_risk_mean": percent_str
}