Yair - added error handling for NA y_train labels

model_trainer.py

@@ -23,7 +23,17 @@ def train_models(X_train, y_train, categorical_columns):
         models["XGBoost"] = xgb
         print(f"✅ XGBoost trained in {time.time() - start_time:.2f} sec")
     else:
-        print("⚠ XGBoost training skipped due to invalid labels!")
+        x_train_xgboost = X_train[~y_train.isna()]
+        y_train_xgboost = y_train.dropna()
+        if set(y_train_xgboost.unique()) <= {0, 1}:
+            start_time = time.time()
+            xgb = XGBClassifier(**XGB_PARAMS)
+            xgb.fit(x_train_xgboost, y_train_xgboost)
+            models["XGBoost"] = xgb
+            print(f"✅ XGBoost trained in {time.time() - start_time:.2f} sec")
+        else:
+            models["XGBoost"] = None
+            print("⚠ XGBoost training skipped due to invalid labels!")
 
     # Train RandomForest
     start_time = time.time()