Spaces:
Sleeping
Sleeping
| """MLflow model registry helpers.""" | |
| import mlflow | |
| import mlflow.pyfunc | |
| import mlflow.sklearn | |
| import mlflow.xgboost | |
| import pandas as pd | |
| from mlflow.tracking import MlflowClient | |
| from sqlalchemy import Engine, text | |
| def register_model( | |
| run_id: str, | |
| model_name: str, | |
| model_uri_suffix: str, | |
| stage: str = "Production", | |
| ) -> None: | |
| """Register a logged model and transition it to the given stage.""" | |
| client = MlflowClient() | |
| model_uri = f"runs:/{run_id}/{model_uri_suffix}" | |
| # Create registered model if it doesn't exist | |
| try: | |
| client.create_registered_model(model_name) | |
| except Exception: | |
| pass # Already exists | |
| version = client.create_model_version( | |
| name=model_name, | |
| source=model_uri, | |
| run_id=run_id, | |
| ) | |
| client.transition_model_version_stage( | |
| name=model_name, | |
| version=version.version, | |
| stage=stage, | |
| archive_existing_versions=True, | |
| ) | |
| print(f" Registered '{model_name}' version {version.version} → {stage}") | |
| def load_production_model(model_name: str): | |
| """Load the Production-stage model from the MLflow registry.""" | |
| return mlflow.pyfunc.load_model(f"models:/{model_name}/Production") | |
| def write_scores_to_warehouse( | |
| engine: Engine, | |
| scores_df: pd.DataFrame, | |
| ) -> int: | |
| """ | |
| Bulk-update dim_customers with ML-derived scores. | |
| scores_df must have column 'customer_unique_id' plus any subset of: | |
| churn_probability, segment_label, rfm_segment, clv_score. | |
| Uses a single connection so temp table is visible across all statements. | |
| """ | |
| possible_cols = ["customer_unique_id", "churn_probability", "segment_label", "rfm_segment", "clv_score"] | |
| cols = [c for c in possible_cols if c in scores_df.columns] | |
| df = scores_df[cols].drop_duplicates("customer_unique_id").copy() | |
| with engine.begin() as conn: | |
| conn.execute(text("DROP TABLE IF EXISTS _scores_staging")) | |
| conn.execute(text(""" | |
| CREATE TEMP TABLE _scores_staging ( | |
| customer_unique_id TEXT, | |
| churn_probability NUMERIC(5,4), | |
| segment_label INTEGER, | |
| rfm_segment TEXT, | |
| clv_score NUMERIC(10,4) | |
| ) | |
| """)) | |
| # Insert using pandas within same connection | |
| ncols = len(cols) | |
| chunk = max(1, 65535 // ncols) | |
| df.to_sql( | |
| "_scores_staging", | |
| con=conn, | |
| if_exists="append", | |
| index=False, | |
| method="multi", | |
| chunksize=chunk, | |
| ) | |
| conn.execute(text(""" | |
| UPDATE warehouse.dim_customers dc | |
| SET | |
| churn_probability = COALESCE(s.churn_probability, dc.churn_probability), | |
| segment_label = COALESCE(s.segment_label, dc.segment_label), | |
| rfm_segment = COALESCE(s.rfm_segment, dc.rfm_segment), | |
| clv_score = COALESCE(s.clv_score, dc.clv_score), | |
| _updated_at = now() | |
| FROM _scores_staging s | |
| WHERE dc.customer_unique_id = s.customer_unique_id | |
| """)) | |
| result = conn.execute(text("SELECT COUNT(*) FROM _scores_staging")) | |
| return result.scalar() | |