Update app.py

app.py

@@ -258,6 +258,8 @@ async def validate_model(
         if os.path.exists(file_path):
             os.remove(file_path)
 
+from predict_utils import run_prediction
+
 @app.post("/v1/xgb/predict")
 async def predict_model(
     file: UploadFile = File(...),
@@ -267,61 +269,15 @@ async def predict_model(
         raise HTTPException(status_code=400, detail="Only CSV files are allowed")
 
     try:
-        # Save uploaded CSV
         file_path = UPLOAD_DIR / file.filename
         with file_path.open("wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
 
-        # Load CSV into DataFrame
         data_df = pd.read_csv(file_path)
         if TEXT_COLUMN not in data_df.columns:
             raise HTTPException(status_code=400, detail=f"Missing required text column: {TEXT_COLUMN}")
 
-        # Load vectorizer
-        tfidf = joblib.load(TFIDF_PATH)
-
-        # Load label encoders
-        label_encoders = joblib.load(ENCODERS_PATH)
-        logger.info(f"Loaded encoders: {list(label_encoders.keys())}")
-
-        # Load model
-        model = TfidfXGBoost(label_encoders)
-        model.load_model(model_name)
-
-        # Vectorize text
-        X_text = data_df[TEXT_COLUMN]
-        X_vec = tfidf.transform(X_text)
-
-        # Make predictions
-        y_pred_list = model.predict(X_vec)
-        all_probs = model.predict_proba(X_vec)
-
-        logger.info(f"y_pred_list: {len(y_pred_list)}, all_probs: {len(all_probs)}")
-
-        # Build prediction response
-        predictions = []
-        for i, col in enumerate(LABEL_COLUMNS):
-            if col not in label_encoders:
-                logger.warning(f"Missing encoder for column: {col}")
-                continue
-
-            label_encoder = label_encoders[col]
-            try:
-                pred_labels = label_encoder.inverse_transform(y_pred_list[i])
-            except Exception as e:
-                logger.error(f"Error in inverse_transform for {col}: {str(e)}")
-                raise HTTPException(status_code=500, detail=f"Label decoding failed for {col}: {str(e)}")
-
-            probs = all_probs[i]
-
-            for pred, prob in zip(pred_labels, probs):
-                predictions.append({
-                    "field": col,
-                    "predicted_label": pred,
-                    "probabilities": {
-                        label: float(p) for label, p in zip(label_encoder.classes_, prob)
-                    }
-                })
+        predictions = run_prediction(data_df, model_name)
 
         return {
             "message": "Prediction completed successfully",