Spaces:
Sleeping
Sleeping
File size: 8,426 Bytes
db4ba8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 | """
TradeFlow AI — XGBoost Rejection Predictor (Phase 3, Step 3.2)
PRD §13 — Trains on submission_outcomes table.
Features: CRS pillars, HS confidence, company historical rate, duty value.
Target: binary (0=accepted, 1=rejected).
The model artifact is stored in MinIO at:
models/rejection_predictor/model_v{version}.joblib
"""
from __future__ import annotations
import io
from pathlib import Path
from typing import Any
import structlog
try:
import numpy as np
except Exception: # pragma: no cover - optional in lightweight test environments
np = None
try:
import joblib
except Exception: # pragma: no cover - optional in lightweight test environments
joblib = None
try:
from sklearn.pipeline import Pipeline
except Exception: # pragma: no cover - optional in lightweight test environments
Pipeline = object
from ..config import settings
log = structlog.get_logger()
FEATURE_NAMES = [
"doc_quality_score",
"completeness_score",
"consistency_score",
"historical_rate",
"hs_confidence",
"cif_value_usd",
"package_count",
"gross_weight_kg",
]
MODEL_LOCAL_PATH = Path("/tmp/rejection_predictor.joblib")
class RejectionPredictor:
"""
XGBoost binary classifier for customs declaration rejection prediction.
"""
def __init__(self) -> None:
self._pipeline: Pipeline | None = None
self._model_version: str = "0"
self._auc: float | None = None
self._load_attempted = False
def _load_from_minio(self) -> bool:
"""Download latest model from MinIO."""
if joblib is None:
log.warning("joblib is not installed; using heuristic predictor fallback")
return False
try:
import boto3
from botocore.client import Config
s3 = boto3.client(
"s3",
endpoint_url=f"http://{settings.MINIO_ENDPOINT}",
aws_access_key_id=settings.MINIO_ACCESS_KEY,
aws_secret_access_key=settings.MINIO_SECRET_KEY,
config=Config(signature_version="s3v4"),
region_name="us-east-1",
)
# List versions, pick latest
response = s3.list_objects_v2(
Bucket=settings.STORAGE_BUCKET_NAME,
Prefix="models/rejection_predictor/",
)
objects = sorted(
response.get("Contents", []),
key=lambda x: x["LastModified"],
reverse=True,
)
if not objects:
log.warning("No model found in MinIO — using heuristic fallback")
return False
latest_key = objects[0]["Key"]
self._model_version = latest_key.split("model_v")[1].replace(".joblib", "")
buf = io.BytesIO()
s3.download_fileobj(settings.STORAGE_BUCKET_NAME, latest_key, buf)
buf.seek(0)
self._pipeline = joblib.load(buf)
log.info("Model loaded from MinIO", version=self._model_version, key=latest_key)
return True
except Exception as exc:
log.error("Failed to load model from MinIO", error=str(exc))
return False
def load(self) -> None:
"""Load model — try MinIO, fallback to local cache."""
self._load_attempted = True
if not self._load_from_minio():
if MODEL_LOCAL_PATH.exists():
if joblib is None:
log.warning("Local model exists but joblib is unavailable")
return
self._pipeline = joblib.load(MODEL_LOCAL_PATH)
log.info("Model loaded from local cache")
else:
log.warning("No model available — predictions will use heuristic")
def predict_proba(self, features: dict) -> float:
"""
Returns probability of rejection (0.0–1.0).
If model is not loaded, falls back to heuristic.
"""
if np is not None:
feature_vec = np.array(
[[features.get(f, 0.0) for f in FEATURE_NAMES]], dtype=np.float32
)
else:
feature_vec = [[features.get(f, 0.0) for f in FEATURE_NAMES]]
if self._pipeline is None and not self._load_attempted:
self.load()
if self._pipeline is not None:
try:
prob = float(self._pipeline.predict_proba(feature_vec)[0][1])
log.debug("XGBoost prediction", prob=prob, version=self._model_version)
return prob
except Exception as exc:
log.error("XGBoost inference error", error=str(exc))
# Heuristic fallback: 1 - average of available scores
known = [v for k, v in features.items() if "score" in k or "confidence" in k]
return round(1.0 - (sum(known) / len(known)), 4) if known else 0.5
def train_and_upload(
self,
X: Any,
y: Any,
version: str,
*,
current_auc: float | None = None,
) -> dict:
"""
Train a new XGBoost model and upload to MinIO.
Called by the Celery `retrain_predictor` task when ≥100 new samples.
"""
if np is None:
raise RuntimeError("numpy is required to train the XGBoost model")
if joblib is None:
raise RuntimeError("joblib is required to train and persist the XGBoost model")
import xgboost as xgb
from sklearn.metrics import roc_auc_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
if len(set(y.tolist())) < 2:
raise ValueError("Cannot train rejection predictor with only one outcome class")
pipeline = Pipeline([
("scaler", StandardScaler()),
("xgb", xgb.XGBClassifier(
n_estimators=300,
max_depth=6,
learning_rate=0.05,
subsample=0.8,
colsample_bytree=0.8,
use_label_encoder=False,
eval_metric="logloss",
random_state=42,
)),
])
stratify = y if min(np.bincount(y.astype(int))) >= 2 else None
X_train, X_valid, y_train, y_valid = train_test_split(
X,
y,
test_size=0.2,
random_state=42,
stratify=stratify,
)
pipeline.fit(X_train, y_train)
valid_prob = pipeline.predict_proba(X_valid)[:, 1]
auc = float(roc_auc_score(y_valid, valid_prob)) if len(set(y_valid.tolist())) > 1 else 0.5
baseline_auc = current_auc if current_auc is not None else self._auc
if baseline_auc is not None and auc < (baseline_auc - settings.RETRAIN_MAX_AUC_DROP):
log.warning(
"Rejected candidate model because AUC regressed",
candidate_auc=auc,
baseline_auc=baseline_auc,
)
return {"promoted": False, "auc": auc, "baseline_auc": baseline_auc}
buf = io.BytesIO()
joblib.dump(pipeline, buf)
buf.seek(0)
# Upload to MinIO
try:
import boto3
from botocore.client import Config
s3 = boto3.client(
"s3",
endpoint_url=f"http://{settings.MINIO_ENDPOINT}",
aws_access_key_id=settings.MINIO_ACCESS_KEY,
aws_secret_access_key=settings.MINIO_SECRET_KEY,
config=Config(signature_version="s3v4"),
region_name="us-east-1",
)
key = f"models/rejection_predictor/model_v{version}.joblib"
s3.upload_fileobj(buf, settings.STORAGE_BUCKET_NAME, key)
log.info("New model uploaded to MinIO", version=version, key=key)
self._pipeline = pipeline
self._model_version = version
self._auc = auc
return {"promoted": True, "auc": auc, "version": version, "key": key}
except Exception as exc:
log.error("Failed to upload model to MinIO", error=str(exc))
raise
# ── Singleton — loaded once per Celery worker startup ─────────────────────────
rejection_predictor = RejectionPredictor()
|