Add ONNX export utilities, pipeline parser, and PyCaret integration

- Implemented `onnx_utils.py` for exporting models to ONNX format and loading ONNX sessions.
- Created `pipeline_parser.py` to infer active steps in AutoML pipelines from logs for various frameworks.
- Developed `pycaret_utils.py` to run PyCaret experiments, including model comparison, tuning, and ONNX export.
- Introduced `training_worker.py` to manage training threads, capturing logs and handling graceful cancellation.
- Added `xai_utils.py` for generating SHAP explanations and occlusion saliency maps for model interpretability.

.dvc/.gitignore

@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache

src/autogluon_utils.py

@@ -3,21 +3,56 @@ import pandas as pd
 import mlflow
 import shutil
 import logging
+import time
+import threading
 from src.mlflow_utils import safe_set_experiment
+from src.onnx_utils import export_to_onnx
 
 logger = logging.getLogger(__name__)
 
 def train_model(train_data: pd.DataFrame, target: str, run_name: str, 
                 valid_data: pd.DataFrame = None, test_data: pd.DataFrame = None, 
-                time_limit: int = 60, presets: str = 'medium_quality', seed: int = 42, cv_folds: int = 0):
+                time_limit: int = 60, presets: str = 'medium_quality', seed: int = 42, cv_folds: int = 0,
+                stop_event=None, task_type: str = "Classification", telemetry_queue=None):
     """
     Trains an AutoGluon model and logs results to MLflow using generic artifact logging.
+    Supports both Tabular data and Computer Vision tasks (via MultiModalPredictor).
     """
-    from autogluon.tabular import TabularPredictor
+    is_cv_task = task_type and task_type.startswith("Computer Vision")
+    is_segmentation = task_type == "Computer Vision - Image Segmentation"
+    is_multilabel = task_type == "Computer Vision - Multi-Label Classification"
+    
+    if is_cv_task:
+        from autogluon.multimodal import MultiModalPredictor
+        
+        def build_image_df(path_df):
+            if path_df is None or "Image_Directory" not in path_df.columns:
+                return path_df
+            img_dir = path_df.iloc[0]["Image_Directory"]
+            data = []
+            for root, _, files in os.walk(img_dir):
+                label = os.path.basename(root)
+                for file in files:
+                    if file.lower().endswith(('.png', '.jpg', '.jpeg')):
+                        data.append({"image": os.path.join(root, file), target: label})
+            return pd.DataFrame(data)
+
+        train_data = build_image_df(train_data)
+        valid_data = build_image_df(valid_data)
+        test_data = build_image_df(test_data)
+    else:
+        from autogluon.tabular import TabularPredictor
     
     safe_set_experiment("AutoGluon_Experiments")
     
-    with mlflow.start_run(run_name=run_name) as run:
+    # Ensure no leaked runs in this thread
+    try:
+        if mlflow.active_run():
+            mlflow.end_run()
+    except:
+        pass
+
+    with mlflow.start_run(run_name=run_name, nested=True) as run:
         # Data cleaning: drop rows where target is NaN
         train_data = train_data.dropna(subset=[target])
         
@@ -44,40 +79,140 @@ def train_model(train_data: pd.DataFrame, target: str, run_name: str,
             test_data = test_data.dropna(subset=[target])
             mlflow.log_param("has_test_data", True)
             
-        # Train model
-        fit_args = {
-            "train_data": train_data,
-            "time_limit": time_limit, 
-            "presets": presets
-        }
-        if cv_folds > 0:
-            fit_args["num_bag_folds"] = cv_folds
+        if is_cv_task:
+            mm_fit_args = {"train_data": train_data, "time_limit": time_limit}
+            if valid_data is not None:
+                mm_fit_args["tuning_data"] = valid_data
             
-        if valid_data is not None and cv_folds == 0:
-            fit_args["tuning_data"] = valid_data
+            problem_type = None
+            if is_segmentation:
+                problem_type = "semantic_segmentation"
+            elif task_type == "Computer Vision - Object Detection":
+                problem_type = "object_detection"
+                
+            mm_presets = "high_quality" if presets in ["best_quality", "high_quality"] else "medium_quality"
+            predictor = MultiModalPredictor(label=target, problem_type=problem_type, path=model_path).fit(**mm_fit_args, presets=mm_presets)
+        else:
+            fit_args = {
+                "train_data": train_data,
+                "time_limit": time_limit, 
+                "presets": presets
+            }
+            if cv_folds > 0:
+                fit_args["num_bag_folds"] = cv_folds
+                
+            if valid_data is not None:
+                fit_args["tuning_data"] = valid_data
+                # If bagging is enabled (manually or by presets), we must set use_bag_holdout=True to use separate tuning_data
+                if cv_folds > 0 or presets in ["best_quality", "high_quality"]:
+                    fit_args["use_bag_holdout"] = True
+                
+            if is_multilabel:
+                fit_args["problem_type"] = "multiclass"
+                mlflow.log_param("is_multilabel", True)
+                
+            # Streaming updates thread
+            def _push_ag_telemetry():
+                while not (stop_event and stop_event.is_set()):
+                    try:
+                        if os.path.exists(model_path):
+                            # AutoGluon sometimes locks the file, so we try-except
+                            from autogluon.tabular import TabularPredictor
+                            try:
+                                temp_predictor = TabularPredictor.load(path=model_path)
+                                lb = temp_predictor.leaderboard(silent=True)
+                                if len(lb) > 0:
+                                    best_model = lb.iloc[0]['model']
+                                    best_score = lb.iloc[0]['score_val']
+                                    if telemetry_queue:
+                                        telemetry_queue.put({
+                                            "status": "running",
+                                            "models_trained": len(lb),
+                                            "best_model": best_model,
+                                            "best_value": best_score,
+                                            "leaderboard_preview": lb.head(5).to_dict(orient='records')
+                                        })
+                            except:
+                                pass
+                    except:
+                        pass
+                    time.sleep(10)
             
-        predictor = TabularPredictor(label=target, path=model_path).fit(**fit_args)
+            if telemetry_queue:
+                t_telemetry = threading.Thread(target=_push_ag_telemetry, daemon=True)
+                t_telemetry.start()
+
+            predictor = TabularPredictor(label=target, path=model_path).fit(**fit_args)
+        
+        # Check if cancelled before continuing
+        if stop_event and stop_event.is_set():
+            raise StopIteration("Training cancelled by user")
         
-        # Log metrics (leaderboard)
-        # If test_data is provided, leaderboard and scoring will strictly use it,
-        # otherwise fallback to training data
         eval_data = test_data if test_data is not None else (valid_data if valid_data is not None else train_data)
-        leaderboard = predictor.leaderboard(eval_data, silent=True)
-        # Log the best model's score
-        best_model_score = leaderboard.iloc[0]['score_val']
-        mlflow.log_metric("best_model_score", best_model_score)
         
-        # Save leaderboard as artifact
-        leaderboard_path = "leaderboard.csv"
-        leaderboard.to_csv(leaderboard_path, index=False)
-        mlflow.log_artifact(leaderboard_path)
-        if os.path.exists(leaderboard_path):
-            os.remove(leaderboard_path)
+        if is_cv_task:
+            scores = predictor.evaluate(eval_data)
+            best_model_score = scores.get('accuracy', scores.get('roc_auc', 0.0))
+            mlflow.log_metrics(scores)
+            leaderboard_path = "leaderboard.csv"
+            pd.DataFrame([scores]).to_csv(leaderboard_path, index=False)
+        else:
+            leaderboard = predictor.leaderboard(eval_data, silent=True)
+            # Log the best model's score
+            best_model_score = leaderboard.iloc[0]['score_val']
+            mlflow.log_metric("best_model_score", best_model_score)
+            leaderboard_path = "leaderboard.csv"
+            leaderboard.to_csv(leaderboard_path, index=False)
+        try:
+            mlflow.log_artifact(leaderboard_path)
+        except Exception as e:
+            logger.warning(f"Failed to log leaderboard artifact: {e}")
+        finally:
+            if os.path.exists(leaderboard_path):
+                os.remove(leaderboard_path)
         
         # Log AutoGluon model directory as a generic artifact
-        # This avoids all ModuleNotFoundError issues with mlflow.autogluon
-        mlflow.log_artifacts(model_path, artifact_path="model")
-        mlflow.log_param("model_type", "autogluon")
+        # We use a try-except here because disk space issues frequently occur during artifact copy
+        try:
+            mlflow.log_artifacts(model_path, artifact_path="model")
+            mlflow.log_param("model_type", "autogluon")
+            
+            # ONNX Export (Best effort for Tabular)
+            if not is_cv_task:
+                try:
+                    onnx_path = os.path.join("models", f"ag_{run_name}.onnx")
+                    # AutoGluon Tabular supports ONNX export for some models
+                    # This might require specific dependencies or AG version
+                    # We call our utility which handles AG logic
+                    export_to_onnx(predictor, "autogluon", target, onnx_path, input_sample=train_data[:1])
+                    mlflow.log_artifact(onnx_path, artifact_path="model")
+                except Exception as e:
+                    logger.warning(f"Failed to export AutoGluon model to ONNX: {e}")
+
+            logger.info(f"AutoGluon artifacts logged successfully for {run_name}")
+            
+            # CRITICAL: Delete local model folder after successful MLflow logging to save disk space
+            # Only do this if it was logged successfully to the tracking server/local mlruns
+            if os.path.exists(model_path):
+                shutil.rmtree(model_path)
+                logger.info(f"Cleaned up local model folder: {model_path}")
+        except Exception as e:
+            logger.error(f"Failed to log model artifacts to MLflow (likely disk space): {e}")
+            # Do NOT delete model_path here so the user can potentially recover it manually
+            # if the MLflow log failed.
+        
+        # Generate and log consumption code sample
+        try:
+            from src.code_gen_utils import generate_consumption_code
+            code_sample = generate_consumption_code("autogluon", run.info.run_id, target)
+            code_path = "consumption_sample.py"
+            with open(code_path, "w") as f:
+                f.write(code_sample)
+            mlflow.log_artifact(code_path)
+            if os.path.exists(code_path):
+                os.remove(code_path)
+        except Exception as e:
+            logger.warning(f"Failed to generate consumption code: {e}")
         
         return predictor, run.info.run_id
 

src/autokeras_utils.py

@@ -0,0 +1,147 @@
+import os
+import shutil
+import time
+import pandas as pd
+import numpy as np
+import mlflow
+import logging
+from src.mlflow_utils import safe_set_experiment
+
+logger = logging.getLogger(__name__)
+
+def run_autokeras_experiment(train_data: pd.DataFrame, target: str, run_name: str, 
+                            valid_data: pd.DataFrame = None, task_type: str = "Computer Vision - Image Classification",
+                            time_limit: int = 60, stop_event=None, log_queue=None):
+    """
+    Trains an AutoKeras model for Image tasks.
+    train_data contains a dataframe with 'Image_Directory' pointing to the dataset path.
+    """
+    safe_set_experiment("AutoKeras_Experiments")
+    
+    try:
+        import autokeras as ak
+        import tensorflow as tf
+    except ImportError:
+        raise ImportError("AutoKeras or TensorFlow not installed. Please install them to use AutoKeras.")
+
+    try:
+        if mlflow.active_run():
+            mlflow.end_run()
+    except Exception:
+        pass
+
+    def qlog(msg):
+        if log_queue:
+            log_queue.put(msg)
+        logger.info(msg)
+
+    with mlflow.start_run(run_name=run_name, nested=True) as run:
+        mlflow.log_param("framework", "autokeras")
+        mlflow.log_param("task_type", task_type)
+        mlflow.log_param("time_limit", time_limit)
+        
+        if "Image_Directory" not in train_data.columns:
+            raise ValueError("AutoKeras requires 'Image_Directory' in the training payload for CV tasks.")
+            
+        img_dir = train_data.iloc[0]["Image_Directory"]
+        qlog(f"Scanning image directory: {img_dir}")
+        
+        # We need to construct tf.data.Dataset from directory
+        # Since AutoKeras ImageClassifier accepts tf.data.Dataset
+        batch_size = 32
+        
+        train_ds = tf.keras.utils.image_dataset_from_directory(
+            img_dir,
+            validation_split=0.2 if valid_data is None else None,
+            subset="training" if valid_data is None else None,
+            seed=42,
+            image_size=(256, 256),
+            batch_size=batch_size
+        )
+        
+        if valid_data is None:
+            val_ds = tf.keras.utils.image_dataset_from_directory(
+                img_dir,
+                validation_split=0.2,
+                subset="validation",
+                seed=42,
+                image_size=(256, 256),
+                batch_size=batch_size
+            )
+        else:
+            val_img_dir = valid_data.iloc[0]["Image_Directory"]
+            val_ds = tf.keras.utils.image_dataset_from_directory(
+                val_img_dir,
+                seed=42,
+                image_size=(256, 256),
+                batch_size=batch_size
+            )
+            
+        mlflow.log_param("num_classes", len(train_ds.class_names))
+        
+        model_path = os.path.join("models", run_name)
+        if os.path.exists(model_path):
+            shutil.rmtree(model_path)
+            
+        qlog("Starting AutoKeras topology search...")
+        
+        # Estimate max trials based on time_limit pseudo translation (1 trial ~ 100s for small data)
+        max_trials = max(1, time_limit // 100)
+        
+        def dataset_to_numpy(ds):
+            x_all, y_all = [], []
+            for x, y in ds:
+                x_all.append(x.numpy())
+                y_all.append(y.numpy())
+            if not x_all: return None, None
+            return np.concatenate(x_all, axis=0), np.concatenate(y_all, axis=0)
+            
+        x_train, y_train = dataset_to_numpy(train_ds)
+        
+        x_val, y_val = None, None
+        if val_ds:
+            x_val, y_val = dataset_to_numpy(val_ds)
+            
+        if task_type == "Computer Vision - Image Classification":
+            clf = ak.ImageClassifier(overwrite=True, max_trials=max_trials, directory=model_path)
+            if val_ds:
+                clf.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=5) # Default short epoch
+            else:
+                clf.fit(x_train, y_train, epochs=5)
+        elif task_type == "Computer Vision - Multi-Label Classification":
+            clf = ak.ImageClassifier(overwrite=True, max_trials=max_trials, directory=model_path, multi_label=True)
+            if val_ds:
+                clf.fit(x_train, y_train, validation_data=(x_val, y_val), epochs=5) # Default short epoch
+            else:
+                clf.fit(x_train, y_train, epochs=5)
+        else:
+            # We don't natively support bounding boxes or segmentation masks right now without specific parser
+            raise NotImplementedError(f"AutoKeras task '{task_type}' requires labels not inherently present in the directory structure or is unsupported by AutoKeras basic API.")
+            
+        if stop_event and stop_event.is_set():
+            raise StopIteration("Training cancelled by user")
+
+        qlog("Evaluating best model...")
+        loss, accuracy = clf.evaluate(val_ds)
+        mlflow.log_metric("val_loss", loss)
+        mlflow.log_metric("val_accuracy", accuracy)
+
+        qlog("Saving and logging artifacts...")
+        export_path = os.path.join(model_path, "best_model")
+        
+        try:
+            model = clf.export_model()
+            model.save(export_path, save_format="tf")
+            mlflow.log_artifacts(export_path, artifact_path="model")
+            mlflow.log_param("model_type", "autokeras")
+            qlog("AutoKeras artifacts logged successfully.")
+        except Exception as e:
+            qlog(f"Warning: Model export failed: {e}")
+
+        # Return a dictionary of useful data for UI
+        return {
+            "run_id": run.info.run_id,
+            "type": "autokeras",
+            # Can't pass TF model across processes easily via queues, so we pass None
+            "predictor": None
+        }

src/code_gen_utils.py

@@ -0,0 +1,422 @@
+import os
+
+
+def generate_consumption_code(model_type: str, run_id: str, target_column: str) -> str:
+    """
+    Generates a Python code snippet to load and run predictions with the trained model.
+    Supports: autogluon, flaml, h2o, tpot, pycaret, lale.
+    """
+    try:
+        client = mlflow.tracking.MlflowClient()
+        run = client.get_run(run_id)
+        task_type = run.data.params.get("task_type", "Classification")
+    except Exception:
+        task_type = "Classification"
+
+    base_code = f"""# Sample code to consume the trained model
+# Run ID: {run_id}
+# Model Type: {model_type}
+# Task Type: {task_type}
+
+import os
+import pandas as pd
+import mlflow
+"""
+
+    if model_type == "autogluon":
+        return base_code + f"""
+from autogluon.tabular import TabularPredictor
+
+# 1. Download model from MLflow
+local_path = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+# 2. Load model
+predictor = TabularPredictor.load(local_path)
+
+# 3. Predict
+# data = pd.read_csv("your_data.csv")
+# predictions = predictor.predict(data)
+# print(predictions)
+"""
+
+    elif model_type == "flaml":
+        return base_code + f"""
+import pickle
+
+# 1. Download model from MLflow
+local_path = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+# 2. Load the .pkl file
+model = None
+for root, dirs, files in os.walk(local_path):
+    for f in files:
+        if f.endswith(".pkl"):
+            with open(os.path.join(root, f), "rb") as fh:
+                model = pickle.load(fh)
+            break
+
+if model is None:
+    raise FileNotFoundError("Model .pkl not found in artifacts.")
+
+# 3. Predict
+# data = pd.read_csv("your_data.csv")
+# predictions = model.predict(data)
+# print(predictions)
+"""
+
+    elif model_type == "h2o":
+        return base_code + f"""
+import h2o
+
+# 1. Initialize H2O
+h2o.init()
+
+# 2. Download model from MLflow
+local_path = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+# 3. Load the H2O model
+model = None
+for root, dirs, files in os.walk(local_path):
+    for f in files:
+        if f.endswith(".zip") or "." not in f:
+            model = h2o.load_model(os.path.join(root, f))
+            break
+
+# 4. Predict
+# h2o_frame = h2o.H2OFrame(pd.read_csv("your_data.csv"))
+# predictions = model.predict(h2o_frame)
+# print(predictions.as_data_frame())
+"""
+
+    elif model_type == "tpot":
+        return base_code + f"""
+import mlflow.sklearn
+
+# 1. Load model directly from MLflow
+model = mlflow.sklearn.load_model("runs:/{run_id}/model")
+
+# 2. Predict
+# data = pd.read_csv("your_data.csv")
+# predictions = model.predict(data)
+# print(predictions)
+"""
+
+    elif model_type == "pycaret":
+        if task_type == "Regression":
+            pc_module = "pycaret.regression"
+        elif task_type == "Time Series Forecasting":
+            pc_module = "pycaret.time_series"
+        else:
+            pc_module = "pycaret.classification"
+
+        return base_code + f"""
+import joblib
+from {pc_module} import load_model, predict_model
+
+# 1. Download model artifact from MLflow
+local_path = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+# 2. Find and load the PyCaret .pkl file
+model_path = None
+for root, dirs, files in os.walk(local_path):
+    for f in files:
+        if f.endswith(".pkl"):
+            model_path = os.path.join(root, f).replace(".pkl", "")
+            break
+
+if model_path is None:
+    raise FileNotFoundError("PyCaret model .pkl not found in artifacts.")
+
+model = load_model(model_path)
+
+# 3. Predict
+# data = pd.read_csv("your_data.csv")  # For classification/regression, must NOT contain target column
+# predictions = predict_model(model, data=data)
+# print(predictions)
+"""
+
+    elif model_type == "lale":
+        return base_code + f"""
+import joblib
+import numpy as np
+from sklearn.preprocessing import OrdinalEncoder, LabelEncoder
+
+# 1. Download model artifact from MLflow
+local_path = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+# 2. Find and load the Lale joblib bundle
+bundle = None
+for root, dirs, files in os.walk(local_path):
+    for f in files:
+        if f.endswith(".pkl"):
+            bundle = joblib.load(os.path.join(root, f))
+            break
+
+if bundle is None:
+    raise FileNotFoundError("Lale model .pkl not found in artifacts.")
+
+model        = bundle["model"]
+col_encoders = bundle.get("col_encoders", {{}})
+y_encoder    = bundle.get("y_encoder", None)
+
+# 3. Preprocess and Predict
+# data = pd.read_csv("your_data.csv")  # must NOT contain target column
+#
+# for col, enc in col_encoders.items():
+#     data[col] = enc.transform(data[[col]]).ravel()
+#
+# raw_preds = model.predict(data.values)
+#
+# if y_encoder is not None:
+#     predictions = y_encoder.inverse_transform(raw_preds)
+# else:
+#     predictions = raw_preds
+#
+# print(predictions)
+"""
+
+    else:
+        return base_code + f"""
+# Code generation for '{model_type}' is not explicitly implemented.
+# Try loading via: mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+"""
+
+
+def _load_code_for_deploy(model_type: str, run_id: str) -> str:
+    """Returns the model-loading block used in the FastAPI main.py."""
+    if model_type == "autogluon":
+        return f"""
+from autogluon.tabular import TabularPredictor
+import mlflow
+_local = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+model = TabularPredictor.load(_local)
+
+def _predict(df):
+    return model.predict(df).tolist()
+"""
+    elif model_type == "flaml":
+        return f"""
+import pickle, os, mlflow
+_local = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+model = None
+for root, _, files in os.walk(_local):
+    for f in files:
+        if f.endswith(".pkl"):
+            with open(os.path.join(root, f), "rb") as fh:
+                model = pickle.load(fh)
+            break
+if model is None:
+    raise FileNotFoundError("FLAML model not found.")
+
+def _predict(df):
+    return model.predict(df).tolist()
+"""
+    elif model_type == "h2o":
+        return f"""
+import h2o, os, mlflow
+h2o.init()
+_local = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+model = None
+for root, _, files in os.walk(_local):
+    for f in files:
+        if f.endswith(".zip") or "." not in f:
+            model = h2o.load_model(os.path.join(root, f))
+            break
+
+def _predict(df):
+    hf = h2o.H2OFrame(df)
+    return model.predict(hf).as_data_frame()["predict"].tolist()
+"""
+    elif model_type == "tpot":
+        return f"""
+import mlflow.sklearn
+model = mlflow.sklearn.load_model("runs:/{run_id}/model")
+
+def _predict(df):
+    return model.predict(df).tolist()
+"""
+    elif model_type == "pycaret":
+        return f"""
+import os, mlflow, joblib
+import pandas as pd
+_local = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+
+try:
+    client = mlflow.tracking.MlflowClient()
+    run = client.get_run("{run_id}")
+    task_type = run.data.params.get("task_type", "Classification")
+except Exception:
+    task_type = "Classification"
+
+if task_type == "Regression":
+    from pycaret.regression import load_model, predict_model
+elif task_type == "Time Series Forecasting":
+    from pycaret.time_series import load_model, predict_model
+else:
+    from pycaret.classification import load_model, predict_model
+
+_mpath = None
+for root, _, files in os.walk(_local):
+    for f in files:
+        if f.endswith(".pkl"):
+            _mpath = os.path.join(root, f).replace(".pkl", "")
+            break
+if _mpath is None:
+    raise FileNotFoundError("PyCaret model not found.")
+model = load_model(_mpath)
+
+def _predict(df):
+    preds = predict_model(model, data=df)
+    if task_type == "Classification" and "prediction_label" in preds.columns:
+        return preds["prediction_label"].tolist()
+    else:
+        # For regression or time series, it might return 'prediction_label' or just predictions
+        if "prediction_label" in preds.columns:
+            return preds["prediction_label"].tolist()
+        return preds.iloc[:, 0].tolist()
+"""
+    elif model_type == "lale":
+        return f"""
+import os, mlflow, joblib
+import numpy as np
+_local = mlflow.artifacts.download_artifacts(run_id="{run_id}", artifact_path="model")
+_bundle = None
+for root, _, files in os.walk(_local):
+    for f in files:
+        if f.endswith(".pkl"):
+            _bundle = joblib.load(os.path.join(root, f))
+            break
+if _bundle is None:
+    raise FileNotFoundError("Lale model not found.")
+_model        = _bundle["model"]
+_col_encoders = _bundle.get("col_encoders", {{}})
+_y_encoder    = _bundle.get("y_encoder", None)
+
+def _predict(df):
+    import pandas as _pd
+    df = _pd.DataFrame(df)
+    for col, enc in _col_encoders.items():
+        if col in df.columns:
+            df[col] = enc.transform(df[[col]]).ravel()
+    raw = _model.predict(df.values)
+    if _y_encoder is not None:
+        return _y_encoder.inverse_transform(raw).tolist()
+    return raw.tolist()
+"""
+    else:
+        return """
+model = None
+def _predict(df):
+    return []
+"""
+
+
+def generate_api_deployment(model_type: str, run_id: str, target_column: str, output_dir: str = "deploy") -> str:
+    """
+    Generates a ready-to-use FastAPI + Docker deployment package for the model.
+    Supports: autogluon, flaml, h2o, tpot, pycaret, lale.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+
+    load_code = _load_code_for_deploy(model_type, run_id)
+
+    main_py = f"""from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+import pandas as pd
+import os
+
+app = FastAPI(title="AutoML Generated API - {model_type}", version="1.0")
+
+# --- Model Loading ---
+{load_code}
+# ---------------------
+
+@app.get("/")
+def health():
+    return {{"status": "running", "model": "{model_type}", "run_id": "{run_id}"}}
+
+@app.post("/predict")
+def predict(payload: dict):
+    try:
+        if "data" in payload:
+            df = pd.DataFrame(payload["data"])
+        else:
+            df = pd.DataFrame([payload])
+        return {{"predictions": _predict(df)}}
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+"""
+
+    with open(os.path.join(output_dir, "main.py"), "w", encoding="utf-8") as f:
+        f.write(main_py)
+
+    # requirements.txt
+    base_reqs = """fastapi==0.104.1
+uvicorn==0.24.0
+pydantic==2.5.2
+pandas==2.1.4
+mlflow==2.9.2
+"""
+    extra = {
+        "autogluon": "autogluon==1.0.0\n",
+        "flaml": "flaml==2.1.2\n",
+        "h2o": "h2o==3.44.0.3\n",
+        "tpot": "tpot==0.12.2\nscikit-learn==1.2.2\n",
+        "pycaret": "pycaret==3.3.0\nscikit-learn==1.2.2\nscipy==1.11.4\n",
+        "lale": "lale==0.9.1\nscikit-learn==1.2.2\njoblib\nhyperopt\n",
+    }
+    reqs = base_reqs + extra.get(model_type, "")
+    with open(os.path.join(output_dir, "requirements.txt"), "w", encoding="utf-8") as f:
+        f.write(reqs)
+
+    # Dockerfile
+    dockerfile = f"""FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y \\
+    build-essential libgomp1 libgl1 python3-dev default-jre curl \\
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY main.py .
+
+EXPOSE 8000
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
+"""
+    with open(os.path.join(output_dir, "Dockerfile"), "w", encoding="utf-8") as f:
+        f.write(dockerfile)
+
+    # README
+    readme = f"""# API Deployment — {model_type} (Run: {run_id})
+
+## Local
+```bash
+pip install -r requirements.txt
+python main.py
+```
+
+## Docker
+```bash
+docker build -t ml-api:{run_id[:8]} .
+docker run -p 8000:8000 ml-api:{run_id[:8]}
+```
+
+## Example request
+```json
+POST http://localhost:8000/predict
+{{
+  "data": [{{"feature1": 1.5, "feature2": "value"}}]
+}}
+```
+"""
+    with open(os.path.join(output_dir, "README.md"), "w", encoding="utf-8") as f:
+        f.write(readme)
+
+    return output_dir

src/data_utils.py

@@ -2,21 +2,40 @@ import os
 import subprocess
 import hashlib
 import time
+import sys
 import pandas as pd
+import zipfile
+import shutil
 
-def load_data(file):
+def load_data(file, no_header=False):
     """
     Loads data from an uploaded file (CSV or Excel) or a disk path.
+    If no_header is True, treats the first row as data (no header) and
+    auto-generates column names as col_0, col_1, ...
     """
     is_path = isinstance(file, str)
     filename = file if is_path else file.name
     
+    if os.path.isdir(filename):
+        # For image directories, return a mock DataFrame to avoid crashing the UI
+        # AutoGluon / AutoKeras will use the path string instead of this DataFrame.
+        num_files = sum(len(files) for _, _, files in os.walk(filename))
+        return pd.DataFrame({"Image_Directory": [filename], "Total_Images": [num_files], "Type": ["Computer Vision Dataset"]})
+        
     if filename.endswith('.csv'):
+        if no_header:
+            df = pd.read_csv(file, header=None)
+            df.columns = [f"col_{i}" for i in range(len(df.columns))]
+            return df
         return pd.read_csv(file)
     elif filename.endswith(('.xls', '.xlsx')):
+        if no_header:
+            df = pd.read_excel(file, header=None)
+            df.columns = [f"col_{i}" for i in range(len(df.columns))]
+            return df
         return pd.read_excel(file)
     else:
-        raise ValueError("Unsupported file format. Please use CSV or Excel.")
+        raise ValueError("Unsupported file format. Please use CSV, Excel, or provide a valid image directory.")
 
 def get_data_summary(df):
     """
@@ -90,14 +109,67 @@ def get_data_lake_files():
         return []
     
     files = []
+    # Add tabular files
     for f in os.listdir(data_lake_dir):
         if f.endswith(('.csv', '.xls', '.xlsx')):
             files.append(os.path.join(data_lake_dir, f))
+            
+    # Add image directories
+    images_dir = os.path.join("data_lake", "images")
+    if os.path.exists(images_dir):
+        for d in os.listdir(images_dir):
+            dir_path = os.path.join(images_dir, d)
+            if os.path.isdir(dir_path):
+                files.append(dir_path)
     
-    # Sort files by creation time descending (newest first)
     files.sort(key=lambda x: os.path.getmtime(x), reverse=True)
     return files
 
+def process_image_upload(uploaded_files, dataset_name="image_dataset", is_zip=False):
+    """
+    Processes uploaded images (multiple files or a zip) and stores them in data_lake/images/<dataset_name>.
+    Supports ZIP extraction or direct copying.
+    Returns the path to the dataset directory and a hash.
+    """
+    data_lake_dir = os.path.join("data_lake", "images", dataset_name)
+    os.makedirs(data_lake_dir, exist_ok=True)
+    
+    timestamp = int(time.time())
+    target_dir = f"{data_lake_dir}_{timestamp}"
+    os.makedirs(target_dir, exist_ok=True)
+
+    if is_zip and len(uploaded_files) == 1:
+        # Extract ZIP
+        zip_file = uploaded_files[0]
+        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+            zip_ref.extractall(target_dir)
+    else:
+        # Multiple Image Files
+        for f in uploaded_files:
+            file_path = os.path.join(target_dir, f.name)
+            with open(file_path, "wb") as out_f:
+                out_f.write(f.getbuffer())
+
+    # Add directory to DVC
+    dvc_hash = "unknown_dir_hash"
+    try:
+        init_dvc()
+        subprocess.run(["dvc", "add", target_dir], check=True, capture_output=True)
+        dvc_file_path = target_dir + ".dvc"
+        if os.path.exists(dvc_file_path):
+            with open(dvc_file_path, "r") as f:
+                content = f.read()
+                import re
+                match = re.search(r'md5:\s*([a-fA-F0-9]+)', content)
+                if match:
+                    dvc_hash = match.group(1)
+    except Exception as e:
+        print(f"DVC error on image dir: {e}")
+        # Pseudo hash fallback
+        dvc_hash = hashlib.md5(target_dir.encode()).hexdigest()
+        
+    return target_dir, dvc_hash, dvc_hash[:8]
+
 def get_dvc_hash(file_path):
     """
     Extracts the DVC hash corresponding to a specific file.

src/experiment_manager.py

@@ -0,0 +1,156 @@
+"""
+ExperimentManager: central registry for all training runs.
+Stored as a singleton in st.session_state['exp_manager'].
+"""
+import threading
+import queue
+import time
+import logging
+from dataclasses import dataclass, field
+from typing import Optional, Any
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExperimentEntry:
+    key: str                          # unique slug: "autogluon_1712345678"
+    metadata: dict                    # framework, run_name, config snapshot
+    thread: Optional[threading.Thread] = field(default=None, repr=False)
+    stop_event: threading.Event = field(default_factory=threading.Event, repr=False)
+    log_queue: queue.Queue = field(default_factory=queue.Queue, repr=False)
+    telemetry_queue: queue.Queue = field(default_factory=queue.Queue, repr=False)
+    result_queue: queue.Queue = field(default_factory=queue.Queue, repr=False)
+    status: str = "queued"            # queued | running | completed | failed | cancelled
+    started_at: float = field(default_factory=time.time)
+    finished_at: Optional[float] = None
+    result: Optional[dict] = None     # {predictor, run_id, type, ...} or {error: str}
+    all_logs: list = field(default_factory=list)
+    latest_telemetry: dict = field(default_factory=dict)
+    last_update: float = field(default_factory=time.time)
+
+    def elapsed_str(self) -> str:
+        end = self.finished_at or time.time()
+        secs = int(end - self.started_at)
+        m, s = divmod(secs, 60)
+        return f"{m}m {s:02d}s"
+
+    def status_icon(self) -> str:
+        return {
+            "queued":    "⏳",
+            "running":   "🟢",
+            "completed": "✅",
+            "failed":    "❌",
+            "cancelled": "🚫",
+        }.get(self.status, "❓")
+
+    def drain_logs(self) -> bool:
+        """Pull all pending log lines and telemetry into the entry."""
+        new = False
+        while not self.log_queue.empty():
+            try:
+                line = self.log_queue.get_nowait()
+                self.all_logs.append(line)
+                new = True
+            except queue.Empty:
+                break
+        
+        while not self.telemetry_queue.empty():
+            try:
+                data = self.telemetry_queue.get_nowait()
+                if isinstance(data, dict):
+                    self.latest_telemetry.update(data)
+                    new = True
+            except queue.Empty:
+                break
+
+        if new:
+            self.last_update = time.time()
+        return new
+
+    def check_result(self):
+        """Non-blocking check: pull result from queue if available."""
+        if not self.result_queue.empty():
+            try:
+                res = self.result_queue.get_nowait()
+                self.result = res
+                if res.get("success"):
+                    self.status = "completed"
+                else:
+                    self.status = "failed"
+                self.finished_at = time.time()
+                self.last_update = time.time()
+            except queue.Empty:
+                pass
+
+
+class ExperimentManager:
+    """In-process registry of all AutoML experiments."""
+
+    def __init__(self):
+        self._runs: dict[str, ExperimentEntry] = {}
+        self._lock = threading.Lock()
+
+    def add(self, entry: ExperimentEntry) -> str:
+        with self._lock:
+            self._runs[entry.key] = entry
+        return entry.key
+
+    def cancel(self, key: str):
+        """Request graceful cancellation of a running experiment."""
+        with self._lock:
+            entry = self._runs.get(key)
+            if entry and entry.status == "running":
+                entry.stop_event.set()
+                entry.status = "cancelled"
+                entry.finished_at = time.time()
+                entry.last_update = time.time()
+                logger.info(f"Cancel requested for experiment: {key}")
+
+    def delete(self, key: str):
+        """Remove experiment from registry (only if not actively running)."""
+        with self._lock:
+            entry = self._runs.get(key)
+            if entry and entry.status == "running":
+                # Cancel first
+                entry.stop_event.set()
+                entry.status = "cancelled"
+                entry.finished_at = time.time()
+                entry.last_update = time.time()
+            self._runs.pop(key, None)
+
+    def get(self, key: str) -> Optional[ExperimentEntry]:
+        with self._lock:
+            return self._runs.get(key)
+
+    def get_all(self) -> list[ExperimentEntry]:
+        """Return all experiments newest-first."""
+        with self._lock:
+            entries = list(self._runs.values())
+        return sorted(entries, key=lambda e: e.started_at, reverse=True)
+
+    def has_running(self) -> bool:
+        return any(e.status == "running" for e in self.get_all())
+
+    def refresh_all(self):
+        """Sync status/logs/results for all experiments."""
+        for entry in self.get_all():
+            entry.drain_logs()
+            if entry.status in ("running", "queued"):
+                entry.check_result()
+                # Also check if thread died unexpectedly
+                if getattr(entry, 'thread', None) is not None:
+                    # Defensive check for is_alive
+                    if not entry.thread.is_alive() and entry.status == "running":
+                        if entry.result is None:
+                            entry.status = "failed"
+                            entry.result = {"success": False, "error": "Thread terminated unexpectedly"}
+                        entry.finished_at = time.time()
+                        entry.last_update = time.time()
+
+
+def get_or_create_manager(session_state) -> ExperimentManager:
+    """Get or create the singleton ExperimentManager from Streamlit session state."""
+    if 'exp_manager' not in session_state or not isinstance(session_state.get('exp_manager'), ExperimentManager):
+        session_state['exp_manager'] = ExperimentManager()
+    return session_state['exp_manager']

src/flaml_utils.py

@@ -1,4 +1,5 @@
 import os
+import threading
 import pandas as pd
 import mlflow
 import shutil
@@ -7,12 +8,16 @@ from flaml import AutoML
 import matplotlib.pyplot as plt
 import time
 from src.mlflow_utils import safe_set_experiment
+from src.onnx_utils import export_to_onnx
 
 logger = logging.getLogger(__name__)
 
 def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str, 
                       valid_data: pd.DataFrame = None, test_data: pd.DataFrame = None,
-                      time_budget: int = 60, task: str = 'classification', metric: str = 'auto', estimator_list: list = 'auto', seed: int = 42, cv_folds: int = 0):
+                       time_budget: int = 60, task: str = 'classification', metric: str = 'auto',
+                       estimator_list: list = 'auto', seed: int = 42, cv_folds: int = 0,
+                       n_jobs: int = 1,
+                       stop_event=None, telemetry_queue=None):
     """
     Trains a FLAML model and logs results to MLflow.
     """
@@ -25,7 +30,14 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
     flaml_logger = logging.getLogger('flaml')
     flaml_logger.setLevel(logging.INFO)
     
-    with mlflow.start_run(run_name=run_name) as run:
+    # Ensure no leaked runs in this thread
+    try:
+        if mlflow.active_run():
+            mlflow.end_run()
+    except:
+        pass
+
+    with mlflow.start_run(run_name=run_name, nested=True) as run:
         # Data cleaning: drop rows where target is NaN
         train_data = train_data.dropna(subset=[target])
         logging.info(f"Data ready: {len(train_data)} rows.")
@@ -48,6 +60,7 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
             valid_data = valid_data.dropna(subset=[target])
             X_val = valid_data.drop(columns=[target])
             y_val = valid_data[target]
+
             mlflow.log_param("has_validation_data", True)
             
         if test_data is not None:
@@ -62,15 +75,16 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
         # The 'No low-cost partial config given' message is just an INFO warning from FLAML.
 
         settings = {
-            "time_budget": time_budget,
             "metric": metric,
             "task": task,
             "estimator_list": estimator_list,
             "log_file_name": "flaml.log",
             "seed": seed,
-            "n_jobs": 1,
+            "n_jobs": n_jobs,
             "verbose": 0, # Reduce internal verbosity to avoid pollution, progress goes to flaml.log
         }
+        if time_budget is not None:
+            settings["time_budget"] = time_budget
         
         if cv_folds > 0:
             settings["eval_method"] = "cv"
@@ -80,6 +94,36 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
             settings["X_val"] = X_val
             settings["y_val"] = y_val
         
+        # Start a watcher thread to respect stop_event
+        _cancel_watcher = None
+        if stop_event is not None:
+            def _watch():
+                stop_event.wait()
+                try:
+                    automl._state.time_budget = 0  # Signal FLAML to stop
+                except Exception:
+                    pass
+            _cancel_watcher = threading.Thread(target=_watch, daemon=True)
+            _cancel_watcher.start()
+
+        # Custom callback for telemetry
+        def _telemetry_callback(iter_count, time_used, best_loss, best_config, estimator, trial_id):
+            try:
+                if telemetry_queue:
+                    telemetry_queue.put({
+                        "status": "running",
+                        "iterations": iter_count,
+                        "time_used": time_used,
+                        "best_loss": best_loss,
+                        "best_estimator": str(estimator),
+                        "best_config_preview": str(best_config)[:200]
+                    })
+            except Exception:
+                pass
+
+        if telemetry_queue:
+            settings["callbacks"] = [_telemetry_callback]
+
         # Train model
         logging.info("Executing hyperparameter search (automl.fit)...")
         try:
@@ -90,6 +134,9 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
             if not hasattr(automl, 'best_estimator') or automl.best_estimator is None:
                 raise RuntimeError("FLAML stopped without finding a valid model.")
         
+        if stop_event and stop_event.is_set():
+            raise StopIteration("Training cancelled by user")
+        
         # Log metrics
         if hasattr(automl, 'best_loss'):
             mlflow.log_metric("best_loss", automl.best_loss)
@@ -106,6 +153,29 @@ def train_flaml_model(train_data: pd.DataFrame, target: str, run_name: str,
         mlflow.log_artifact(model_path, artifact_path="model")
         mlflow.log_param("model_type", "flaml")
         
+        # ONNX Export
+        try:
+            onnx_path = os.path.join("models", f"flaml_{run_name}.onnx")
+            # For FLAML, we can often export the underlying best estimator or the AutoML object if it's scikit-learn compatible
+            # We pass X_train[:1] as sample input for shape inference
+            export_to_onnx(automl.model.estimator, "flaml", target, onnx_path, input_sample=X_train[:1])
+            mlflow.log_artifact(onnx_path, artifact_path="model")
+        except Exception as e:
+            logger.warning(f"Failed to export FLAML model to ONNX: {e}")
+        
+        # Generate and log consumption code sample
+        try:
+            from src.code_gen_utils import generate_consumption_code
+            code_sample = generate_consumption_code("flaml", run.info.run_id, target)
+            code_path = "consumption_sample.py"
+            with open(code_path, "w") as f:
+                f.write(code_sample)
+            mlflow.log_artifact(code_path)
+            if os.path.exists(code_path):
+                os.remove(code_path)
+        except Exception as e:
+            logger.warning(f"Failed to generate consumption code: {e}")
+            
         # Log training log as artifact
         if os.path.exists("flaml.log"):
             mlflow.log_artifact("flaml.log")

src/h2o_utils.py

@@ -111,7 +111,8 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
                    valid_data: pd.DataFrame = None, test_data: pd.DataFrame = None,
                    max_runtime_secs: int = 300, max_models: int = 10, 
                    nfolds: int = 3, balance_classes: bool = True, seed: int = 42,
-                   sort_metric: str = "AUTO", exclude_algos: list = None):
+                   sort_metric: str = "AUTO", exclude_algos: list = None,
+                   stop_event=None, telemetry_queue=None):
     """
     Trains H2O AutoML model and registers in MLflow
     """
@@ -125,7 +126,14 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
     h2o_instance = initialize_h2o()
     
     try:
-        with mlflow.start_run(run_name=run_name) as run:
+        # Ensure no leaked runs in this thread
+        try:
+            if mlflow.active_run():
+                mlflow.end_run()
+        except:
+            pass
+
+        with mlflow.start_run(run_name=run_name, nested=True) as run:
             # Prepare data
             h2o_frame, clean_data = prepare_data_for_h2o(train_data, target)
             
@@ -158,6 +166,23 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
                 sort_metric=sort_metric,
                 exclude_algos=exclude_algos or []
             )
+
+            # Watcher thread for graceful cancellation
+            import threading
+            _cancel_watcher = None
+            if stop_event is not None:
+                def _h2o_watch():
+                    stop_event.wait()
+                    try:
+                        import h2o as _h2o
+                        jobs = _h2o.cluster().jobs()
+                        for job in jobs:
+                            if job.status == 'RUNNING':
+                                job.cancel()
+                    except Exception:
+                        pass
+                _cancel_watcher = threading.Thread(target=_h2o_watch, daemon=True)
+                _cancel_watcher.start()
             
             # Prepare test and validation data if present
             h2o_valid = None
@@ -178,14 +203,62 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
             
             # Train model
             logger.info("Starting H2O AutoML training...")
+            import sys
+            # Guard against deep recursion in H2O/Scipy on some datasets
+            sys.setrecursionlimit(max(sys.getrecursionlimit(), 3000))
+
             start_time = time.time()
             train_kwargs = {"x": features, "y": target, "training_frame": h2o_frame}
             if h2o_valid is not None:
                 train_kwargs["validation_frame"] = h2o_valid
             if h2o_test is not None:
                 train_kwargs["leaderboard_frame"] = h2o_test
+            
+            # Streaming updates thread
+            def _push_h2o_telemetry():
+                while aml.leaderboard is None or aml.leaderboard.nrow == 0:
+                    if stop_event and stop_event.is_set(): break
+                    time.sleep(2)
                 
-            aml.train(**train_kwargs)
+                last_row_count = 0
+                while not (stop_event and stop_event.is_set()):
+                    try:
+                        lb = aml.leaderboard
+                        if lb is not None and lb.nrow > last_row_count:
+                            last_row_count = lb.nrow
+                            lb_df = lb.as_data_frame()
+                            best_metric = lb_df.columns[1] if len(lb_df.columns) > 1 else "score"
+                            best_val = lb_df.iloc[0, 1] if len(lb_df) > 0 else 0
+                            
+                            if telemetry_queue:
+                                telemetry_queue.put({
+                                    "status": "running",
+                                    "models_trained": last_row_count,
+                                    "best_metric": best_metric,
+                                    "best_value": best_val,
+                                    "leaderboard_preview": lb_df.head(5).to_dict(orient='records')
+                                })
+                    except Exception:
+                        pass
+                    if training_duration := (time.time() - start_time):
+                         if training_duration > max_runtime_secs and max_runtime_secs > 0: break
+                    time.sleep(5)
+
+            if telemetry_queue:
+                t_telemetry = threading.Thread(target=_push_h2o_telemetry, daemon=True)
+                t_telemetry.start()
+
+            # Fix encoding issue on Windows by disabling H2O progress bar if it causes issues
+            # or wrapping the call. H2O uses ASCII bars if it detects non-tty, but our router
+            # might be confusing it.
+            try:
+                aml.train(**train_kwargs)
+            except UnicodeEncodeError:
+                # Fallback: try with minimal verbosity if encoding fails
+                logger.warning("Encoding error detected, retrying with lower verbosity...")
+                aml.project_name = aml.project_name + "_retry"
+                aml.train(**train_kwargs)
+
             training_duration = time.time() - start_time
             
             logger.info(f"Training completed in {training_duration:.2f} seconds")
@@ -214,51 +287,56 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
             
             # Save leaderboard as metric with safe wrapper
             try:
-                # Check available columns in leaderboard
-                leaderboard_df = None
+                available_metrics = []
+                num_models = 0
+                
                 try:
+                    num_models = leaderboard.nrow
                     leaderboard_df = leaderboard.as_data_frame()
-                    logger.info(f"Available columns: {list(leaderboard_df.columns)}")
+                    available_metrics = [c.lower() for c in leaderboard_df.columns]
+                    logger.info(f"Available leaderboard columns: {list(leaderboard_df.columns)}")
                 except Exception as e:
-                    logger.warning(f"Could not convert leaderboard to DataFrame: {e}")
-                
-                # Try to get the best available metric
+                    logger.warning(f"Metadata extraction failed: {e}")
+                    leaderboard_df = None
+
+                # Search for metrics in preference order
                 best_model_score = 0.0
-                if leaderboard_df is not None and len(leaderboard_df) > 0:
-                    # Search for metrics in preference order
-                    for metric in ['auc', 'logloss', 'rmse', 'mae', 'r2']:
-                        if metric in leaderboard_df.columns:
-                            best_model_score = leaderboard_df.iloc[0][metric]
-                            logger.info(f"Using metric '{metric}': {best_model_score}")
+                found_metric = "none"
+                
+                metric_candidates = ['auc', 'logloss', 'rmse', 'mae', 'r2', 'mse', 'accuracy', 'f1']
+                
+                if leaderboard_df is not None and not leaderboard_df.empty:
+                    # Find column index
+                    col_names_lower = [c.lower() for c in leaderboard_df.columns]
+                    for m_cand in metric_candidates:
+                        if m_cand in col_names_lower:
+                            idx = col_names_lower.index(m_cand)
+                            actual_col = leaderboard_df.columns[idx]
+                            best_model_score = float(leaderboard_df.iloc[0][actual_col])
+                            found_metric = actual_col
+                            logger.info(f"Using metric '{found_metric}': {best_model_score}")
                             break
                     
-                    mlflow.log_metric("total_models_trained", len(leaderboard_df))
-                else:
-                    # Fallback: use the first value in H2O leaderboard
-                    try:
-                        available_columns = leaderboard.columns
-                        logger.info(f"Available H2O columns: {available_columns}")
-                        
-                        # Try accessing first row, first metric col
-                        if len(available_columns) > 0:
-                            first_col = available_columns[0]
-                            best_model_score = leaderboard[0, first_col]
-                            logger.info(f"Using first available column '{first_col}': {best_model_score}")
-                        
-                        mlflow.log_metric("total_models_trained", leaderboard.nrow)
-                    except Exception as e:
-                        logger.warning(f"Could not extract metrics from leaderboard: {e}")
-                        mlflow.log_metric("total_models_trained", 0)
+                    # If still 0 and we have columns, pick the second one (usually the main metric)
+                    if best_model_score == 0.0 and len(leaderboard_df.columns) > 1:
+                        actual_col = leaderboard_df.columns[1]
+                        best_model_score = float(leaderboard_df.iloc[0][actual_col])
+                        found_metric = actual_col
+                        logger.info(f"Fallback to second column '{found_metric}': {best_model_score}")
                 
+                # Log metrics
+                mlflow.log_metric("total_models_trained", float(num_models))
                 mlflow.log_metric("best_model_score", best_model_score)
                 mlflow.log_metric("training_duration", training_duration)
+                if found_metric != "none":
+                    mlflow.set_tag("best_metric_name", found_metric)
                 
             except Exception as e:
                 logger.warning(f"Error processing leaderboard metrics: {e}")
-                # Default fallback
+                # Ultimate fallback
                 mlflow.log_metric("best_model_score", 0.0)
                 mlflow.log_metric("training_duration", training_duration)
-                mlflow.log_metric("total_models_trained", 0)
+                mlflow.log_metric("total_models_trained", 0.0)
             
             # Try saving leaderboard with error handling
             try:
@@ -297,6 +375,19 @@ def train_h2o_model(train_data: pd.DataFrame, target: str, run_name: str,
                 h2o.save_model(best_model, path=temp_model_path)
                 mlflow.log_artifacts(temp_model_path, artifact_path="model")
                 
+                # Generate and log consumption code sample
+                try:
+                    from src.code_gen_utils import generate_consumption_code
+                    code_sample = generate_consumption_code("h2o", run.info.run_id, target)
+                    code_path = "consumption_sample.py"
+                    with open(code_path, "w") as f:
+                        f.write(code_sample)
+                    mlflow.log_artifact(code_path)
+                    if os.path.exists(code_path):
+                        os.remove(code_path)
+                except Exception as e:
+                    logger.warning(f"Failed to generate consumption code: {e}")
+
                 # Clean temp directory
                 import shutil
                 if os.path.exists(temp_model_path):

src/huggingface_utils.py

@@ -0,0 +1,96 @@
+import os
+import logging
+from typing import List, Dict, Any, Optional
+
+logger = logging.getLogger(__name__)
+
+HF_AVAILABLE = None
+
+def _check_hf_availability():
+    global HF_AVAILABLE
+    if HF_AVAILABLE is not None:
+        return HF_AVAILABLE
+    try:
+        import huggingface_hub
+        HF_AVAILABLE = True
+    except Exception as e:
+        logger.warning(f"Hugging Face Hub not available: {e}")
+        HF_AVAILABLE = False
+    return HF_AVAILABLE
+
+class HuggingFaceService:
+    def __init__(self, token: Optional[str] = None):
+        if not _check_hf_availability():
+            self.api = None
+            self.token = None
+            return
+
+        from huggingface_hub import HfApi
+        self.token = token or os.environ.get("HUGGINGFACE_TOKEN")
+        self.api = HfApi(token=self.token) if self.token else HfApi()
+        
+    def authenticate(self, token: str):
+        """Authenticates with the Hugging Face Hub."""
+        if not _check_hf_availability():
+            raise ImportError("Hugging Face Hub library not found.")
+            
+        from huggingface_hub import login, HfApi
+        self.token = token
+        self.api = HfApi(token=token)
+        login(token=token)
+        logger.info("Authenticated with Hugging Face Hub.")
+
+    def list_models(self, query: str = None, author: str = None) -> List[Dict[str, Any]]:
+        """Lists models on the Hub based on search query or author."""
+        if not self.api:
+            return []
+        models = self.api.list_models(search=query, author=author, limit=10)
+        return [{"id": m.id, "author": m.author, "lastModified": m.lastModified} for m in models]
+
+    def upload_model(self, model_path: str, repo_id: str, commit_message: str = "Upload AutoML model", private: bool = True):
+        """Uploads a model file or directory to a HF repository."""
+        if not self.api:
+            raise ImportError("Hugging Face Hub library not found.")
+        if not self.token:
+            raise ValueError("Authentication token is required for upload.")
+            
+        repo_url = self.api.create_repo(repo_id=repo_id, private=private, exist_ok=True)
+        logger.info(f"Hub repository ready: {repo_url}")
+        
+        if os.path.isdir(model_path):
+            self.api.upload_folder(
+                folder_path=model_path,
+                repo_id=repo_id,
+                commit_message=commit_message
+            )
+        else:
+            self.api.upload_file(
+                path_or_fileobj=model_path,
+                path_in_repo=os.path.basename(model_path),
+                repo_id=repo_id,
+                commit_message=commit_message
+            )
+        logger.info(f"Successfully uploaded {model_path} to {repo_id}")
+
+    def download_model(self, repo_id: str, filename: str, local_dir: str = "models/hf_downloads") -> str:
+        """Downloads a specific file from a HF repository."""
+        if not _check_hf_availability():
+            raise ImportError("Hugging Face Hub library not found.")
+            
+        from huggingface_hub import hf_hub_download
+        os.makedirs(local_dir, exist_ok=True)
+        path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=local_dir)
+        logger.info(f"Downloaded {filename} from {repo_id} to {path}")
+        return path
+
+    def consult_model_info(self, repo_id: str) -> Dict[str, Any]:
+        """Gets metadata about a model on the Hub."""
+        if not self.api:
+            return {}
+        info = self.api.model_info(repo_id=repo_id)
+        return {
+            "id": info.id,
+            "tags": info.tags,
+            "pipeline_tag": info.pipeline_tag,
+            "downloads": info.downloads
+        }

src/lale_utils.py

@@ -0,0 +1,198 @@
+import os
+import logging
+import traceback
+import queue
+import time
+import pandas as pd
+import numpy as np
+import joblib
+from typing import Dict, Any, Optional
+
+import mlflow
+
+# Lale core imports
+import lale
+from lale.lib.lale import Hyperopt
+from lale.lib.sklearn import LogisticRegression, RandomForestClassifier
+from lale.lib.sklearn import MinMaxScaler, PCA
+from sklearn.preprocessing import LabelEncoder, OrdinalEncoder
+
+from src.mlflow_utils import safe_set_experiment
+
+
+def _preprocess_for_lale(X: pd.DataFrame, y: pd.Series, task_type: str = "Classification"):
+    """
+    Encode non-numeric features so that sklearn estimators can handle them.
+    Returns (X_encoded, y_encoded, encoders) where encoders can be used for inverse transforms.
+    """
+    X = X.copy()
+
+    # Encode categorical / object columns
+    col_encoders = {}
+    for col in X.columns:
+        if X[col].dtype == object or str(X[col].dtype) == 'category':
+            le = OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=-1)
+            X[col] = le.fit_transform(X[[col]]).ravel()
+            col_encoders[col] = le
+
+    # Fill any remaining NaNs
+    for col in X.columns:
+        if X[col].isna().any():
+            X[col] = X[col].fillna(X[col].median() if pd.api.types.is_numeric_dtype(X[col]) else 0)
+
+    # Encode target if classification (Regression target should remain continuous)
+    y_encoder = None
+    if task_type != "Regression":
+        if y.dtype == object or str(y.dtype) == 'category':
+            y_encoder = LabelEncoder()
+            y = pd.Series(y_encoder.fit_transform(y), name=y.name)
+
+    return X, y, col_encoders, y_encoder
+
+
+def run_lale_experiment(
+    train_df: pd.DataFrame,
+    target_col: str,
+    run_name: str,
+    time_limit: Optional[int],
+    log_queue: queue.Queue,
+    stop_event=None,
+    val_df: Optional[pd.DataFrame] = None,
+    task_type: str = "Classification",
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Run Lale experiment using scikit-learn compatible pipelines via Hyperopt.
+    Handles text/categorical features with automatic encoding.
+    """
+    logger = logging.getLogger("lale")
+    logger.info(f"Starting Lale experiment: {run_name} (Task: {task_type})")
+    logger.info(f"Dataset shape: {train_df.shape}, Target: {target_col}")
+
+    # Drop NaNs on target
+    train_df_c = train_df.dropna(subset=[target_col])
+    X_raw = train_df_c.drop(columns=[target_col])
+    y_raw = train_df_c[target_col]
+
+    # Pre-process: encode categoricals/text for sklearn compatibility
+    logger.info("Step: Encoding categorical/text features...")
+    X, y, col_encoders, y_encoder = _preprocess_for_lale(X_raw, y_raw, task_type)
+    
+    unique_classes_log = ""
+    if task_type != "Regression":
+       unique_classes_log = f" | Classes: {y.unique()[:5].tolist()}"
+    logger.info(f"Features after encoding: {list(X.columns)}{unique_classes_log}")
+
+    # Validate MLflow tracking
+    safe_set_experiment("Multi_AutoML_Project")
+
+    # Always end any dangling run (Hyperopt can leave runs open)
+    try:
+        mlflow.end_run()
+    except Exception:
+        pass
+
+    if stop_event and stop_event.is_set():
+        raise StopIteration("Experiment cancelled before setup.")
+
+    try:
+        with mlflow.start_run(run_name=run_name) as run:
+            run_id = run.info.run_id
+            logger.info(f"MLflow Run ID: {run_id}")
+            mlflow.log_param("model_type", "lale")
+            mlflow.log_param("n_features", X.shape[1])
+            mlflow.log_param("n_samples", X.shape[0])
+            mlflow.log_param("task_type", task_type)
+
+            # 1. Pipeline Definition (only numeric-friendly preprocessors)
+            logger.info("Step: Defining Lale Planned Pipeline...")
+            
+            if task_type == "Regression":
+                from lale.lib.sklearn import LinearRegression, RandomForestRegressor
+                planned_pipeline = (
+                    (MinMaxScaler | PCA) >>
+                    (LinearRegression | RandomForestRegressor)
+                )
+                scoring_metric = "r2"
+            else:
+                planned_pipeline = (
+                    (MinMaxScaler | PCA) >>
+                    (LogisticRegression | RandomForestClassifier)
+                )
+                scoring_metric = "accuracy"
+
+            if stop_event and stop_event.is_set():
+                raise StopIteration("Experiment cancelled before Hyperopt setup.")
+
+            # 2. Hyperparameter Tuning
+            logger.info("Step: Tuning with Hyperopt...")
+            max_evals = 10 if time_limit is None or time_limit >= 300 else 5
+            time_args = {}
+            if time_limit and time_limit > 0:
+                time_args['max_eval_time'] = time_limit
+
+            optimizer = Hyperopt(
+                estimator=planned_pipeline,
+                max_evals=max_evals,
+                cv=3,
+                scoring=scoring_metric,
+                show_progressbar=False,
+                verbose=True,   # show per-trial info so we can debug failures
+                **time_args
+            )
+
+            # 3. Fit Model
+            logger.info(f"Step: Fitting Lale Optimizer (evals={max_evals})...")
+            start_time = time.time()
+            trained_optimizer = optimizer.fit(X.values, y.values)
+
+            if stop_event and stop_event.is_set():
+                raise StopIteration("Experiment cancelled after fitting.")
+
+            best_model = trained_optimizer.get_pipeline()
+
+            # Extract score
+            try:
+                summary = trained_optimizer.summary()
+                best_score = -summary.iloc[0]['loss'] if 'loss' in summary.columns else 0.0
+            except Exception:
+                best_score = 0.0
+
+            elapsed_time = time.time() - start_time
+            logger.info(f"Best Score (CV {scoring_metric}): {best_score:.4f}")
+            logger.info(f"Optimization time: {elapsed_time:.1f}s")
+
+            # 4. Save Model
+            logger.info("Step: Saving model locally...")
+            model_dir = "models"
+            os.makedirs(model_dir, exist_ok=True)
+            model_path = os.path.join(model_dir, f"{run_name}_lale_model.pkl")
+            joblib.dump({"model": best_model, "col_encoders": col_encoders, "y_encoder": y_encoder, "task_type": task_type}, model_path)
+
+            # Log metrics
+            mlflow.log_metric(f"best_cv_{scoring_metric}", best_score)
+            mlflow.log_metric("optimization_time", elapsed_time)
+            mlflow.log_param("max_evals", max_evals)
+            mlflow.log_artifact(model_path, artifact_path="model")
+
+            logger.info("Lale experiment completed successfully.")
+
+            # 5. Prepare return bundle
+            bundle = {"model": best_model, "col_encoders": col_encoders, "y_encoder": y_encoder, "task_type": task_type}
+
+            return {
+                "success": True,
+                "predictor": bundle,
+                "run_id": run_id,
+                "type": "lale",
+                "model_path": model_path,
+                "metrics": {f"best_cv_{scoring_metric}": best_score}
+            }
+
+    except StopIteration as si:
+        logger.warning(f"Cancelled: {si}")
+        raise
+    except Exception as e:
+        logger.error(f"Lale Error: {e}")
+        logger.error(traceback.format_exc())
+        raise e

src/onnx_utils.py

@@ -0,0 +1,105 @@
+import os
+import logging
+from typing import Any, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+# Global flags for availability
+ONNX_AVAILABLE = None
+
+def _check_onnx_availability():
+    global ONNX_AVAILABLE
+    if ONNX_AVAILABLE is not None:
+        return ONNX_AVAILABLE
+    try:
+        import onnx
+        import onnxruntime as ort
+        ONNX_AVAILABLE = True
+    except Exception as e:
+        logger.warning(f"ONNX or ONNXRuntime not available: {e}")
+        ONNX_AVAILABLE = False
+    return ONNX_AVAILABLE
+
+def export_to_onnx(model: Any, model_type: str, target_col: str, output_path: str, input_sample: Optional[Any] = None) -> str:
+    """
+    Exports a trained model to ONNX format.
+    Supports: flaml, pycaret, autogluon (tabular), autokeras (tensorflow).
+    """
+    if not _check_onnx_availability():
+        raise ImportError("ONNX or ONNXRuntime is not available in this environment.")
+    
+    import onnx
+    import pandas as pd
+    import numpy as np
+    
+    logger.info(f"Exporting {model_type} model to ONNX: {output_path}")
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    try:
+        if model_type in ["flaml", "pycaret", "tpot"]:
+            from skl2onnx import to_onnx
+            
+            if input_sample is None:
+                raise ValueError("input_sample is required for scikit-learn based ONNX export")
+            
+            if isinstance(input_sample, pd.DataFrame) and target_col in input_sample.columns:
+                input_sample = input_sample.drop(columns=[target_col])
+            
+            onx = to_onnx(model, input_sample[:1], initial_types=None)
+            with open(output_path, "wb") as f:
+                f.write(onx.SerializeToString())
+
+        elif model_type == "autokeras":
+            import tf2onnx
+            import tensorflow as tf
+            
+            if input_sample is None:
+                raise ValueError("input_sample is required for TensorFlow/AutoKeras ONNX export")
+                
+            input_signature = [tf.TensorSpec([None] + list(input_sample.shape[1:]), tf.float32, name='input')]
+            onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=13)
+            onnx.save_model(onnx_model, output_path)
+
+        elif model_type == "autogluon":
+            try:
+                model.export_onnx(output_path)
+            except AttributeError:
+                logger.warning("AutoGluon model does not support direct export_onnx.")
+                raise NotImplementedError("AutoGluon ONNX export fallback not implemented.")
+
+        else:
+            raise ValueError(f"Unsupported model type for ONNX export: {model_type}")
+
+        logger.info(f"Successfully exported model to {output_path}")
+        return output_path
+
+    except Exception as e:
+        logger.error(f"Failed to export {model_type} model to ONNX: {e}")
+        raise
+
+def load_onnx_session(onnx_path: str):
+    """Loads an ONNX model into an inference session."""
+    if not _check_onnx_availability():
+        raise ImportError("ONNXRuntime is not available.")
+    
+    import onnxruntime as ort
+    if not os.path.exists(onnx_path):
+        raise FileNotFoundError(f"ONNX file not found: {onnx_path}")
+    return ort.InferenceSession(onnx_path)
+
+def predict_onnx(session: Any, df: Any) -> Any:
+    """Runs inference on a DataFrame using an ONNX session."""
+    import numpy as np
+    
+    inputs = {}
+    for node in session.get_inputs():
+        name = node.name
+        if name in df.columns:
+            inputs[name] = df[[name]].values.astype(np.float32)
+        else:
+            if len(session.get_inputs()) == 1:
+                inputs[name] = df.values.astype(np.float32)
+                break
+    
+    outputs = session.run(None, inputs)
+    return outputs[0]

src/pipeline_parser.py

@@ -0,0 +1,328 @@
+"""
+pipeline_parser.py — infer which AutoML pipeline step is active from live logs.
+
+Each framework has a sequence of steps. This module parses log lines
+to determine which step is "done", which is "active", and which is "pending".
+"""
+
+from typing import Optional
+
+# ── Step definitions per framework ───────────────────────────────────────────
+# Each step has:
+#   label       — displayed name
+#   keywords    — log keywords that signal this step has STARTED or is active
+#   done_kw     — log keywords that signal this step is DONE (optional)
+#   description — tooltip / explainer text
+
+_STEPS: dict[str, list[dict]] = {
+    "autogluon": [
+        {
+            "label": "Data Preparation",
+            "icon": "📊",
+            "keywords": ["preprocessing", "converting", "fitting", "loading data", "train_data"],
+            "done_kw": ["beginning automl", "fitting model:"],
+            "description": "Validates and preprocesses the dataset. Handles missing values, categorical encoding and feature types.",
+        },
+        {
+            "label": "Fitting Models",
+            "icon": "🤖",
+            "keywords": ["fitting model:", "training model for", "fitting with cpus"],
+            "done_kw": ["weightedensemble", "autogluon training complete"],
+            "description": "Trains each individual model (LightGBM, XGBoost, CatBoost, RF, etc.) within the time budget.",
+        },
+        {
+            "label": "Stacking / Ensembling",
+            "icon": "🏗️",
+            "keywords": ["weightedensemble", "ensemble weights", "stacking"],
+            "done_kw": ["autogluon training complete"],
+            "description": "Combines the best models using weighted ensembling or multi-layer stacking.",
+        },
+        {
+            "label": "Evaluation",
+            "icon": "📏",
+            "keywords": ["leaderboard", "best model:", "validation score", "score_val"],
+            "done_kw": ["tabularpredictor saved", "best model logged"],
+            "description": "Evaluates all models on the validation set and builds the final leaderboard.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["mlflow", "log_artifacts", "logged successfully", "artifacts logged"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+    "flaml": [
+        {
+            "label": "Data Preparation",
+            "icon": "📊",
+            "keywords": ["data ready", "preprocessing", "starting flaml"],
+            "done_kw": ["executing hyperparameter search"],
+            "description": "Validates the dataset, detects feature types, and prepares inputs for FLAML's optimizer.",
+        },
+        {
+            "label": "Hyperparameter Search",
+            "icon": "🔍",
+            "keywords": ["executing hyperparameter search", "automl.fit", "[flaml.automl", "trial", "best config"],
+            "done_kw": ["search finished"],
+            "description": "FLAML runs a cost-effective search over hyperparameter configurations using Bayesian optimization.",
+        },
+        {
+            "label": "Best Config Selection",
+            "icon": "🏆",
+            "keywords": ["search finished", "best estimator", "best loss", "best final"],
+            "done_kw": ["saving best model"],
+            "description": "Identifies the best-performing estimator and its configuration from the search results.",
+        },
+        {
+            "label": "Model Saving",
+            "icon": "💾",
+            "keywords": ["saving best model", "model_path", "artifact_path"],
+            "done_kw": ["mlflow", "logged successfully"],
+            "description": "Serializes the trained model to disk using pickle.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["mlflow", "log_artifact", "logged successfully"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+    "h2o": [
+        {
+            "label": "H2O Cluster Init",
+            "icon": "🌊",
+            "keywords": ["h2o cluster initialized", "initializing h2o", "h2o init"],
+            "done_kw": ["starting h2o automl"],
+            "description": "Starts the local H2O Java cluster and allocates memory for distributed model training.",
+        },
+        {
+            "label": "Data Preparation",
+            "icon": "📊",
+            "keywords": ["preparing data", "h2oframe", "feature engineering", "asfactor"],
+            "done_kw": ["starting h2o automl training"],
+            "description": "Converts Pandas DataFrames to H2O frames and applies type casting for features/targets.",
+        },
+        {
+            "label": "AutoML Training",
+            "icon": "🤖",
+            "keywords": ["starting h2o automl training", "automl session", "training completed", "aml.train"],
+            "done_kw": ["training completed in"],
+            "description": "H2O trains multiple model families (GBM, XGBoost, GLM, DRF, DeepLearning) and their variants.",
+        },
+        {
+            "label": "Leaderboard & Scoring",
+            "icon": "📏",
+            "keywords": ["top 5 models", "leaderboard", "best model score", "auc", "total_models_trained"],
+            "done_kw": ["model saved at", "log model to mlflow"],
+            "description": "Ranks all trained models and evaluates the leader on the validation/test set.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["mlflow", "log_artifacts", "logged successfully", "artifacts logged"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+    "tpot": [
+        {
+            "label": "Data Preparation",
+            "icon": "📊",
+            "keywords": ["problem type:", "training data shape", "test data shape", "label encoder"],
+            "done_kw": ["starting tpot training"],
+            "description": "Applies feature engineering pipelines: TF-IDF for text, ordinal encoding, and standard scaling.",
+        },
+        {
+            "label": "Pipeline Generation (GA)",
+            "icon": "🧬",
+            "keywords": ["starting tpot training", "generation:", "pipeline score:", "optimizing pipeline"],
+            "done_kw": ["training completed"],
+            "description": "TPOT uses a Genetic Algorithm to evolve and select the best scikit-learn pipeline configurations.",
+        },
+        {
+            "label": "Pipeline Selection",
+            "icon": "🏆",
+            "keywords": ["training completed", "best pipeline", "fitted_pipeline_", "accuracy:", "f1_macro:"],
+            "done_kw": ["pipeline exported"],
+            "description": "Identifies the highest-scoring pipeline from the genetic search as the final model.",
+        },
+        {
+            "label": "Export & Analysis",
+            "icon": "📤",
+            "keywords": ["pipeline exported", "export", "classification report"],
+            "done_kw": ["mlflow"],
+            "description": "Exports the best pipeline as a .py file and generates a classification/regression report.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["mlflow", "tpot automl model", "registered_model_name", "logged successfully"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+    "pycaret": [
+        {
+            "label": "Environment Setup",
+            "icon": "⚙️",
+            "keywords": ["setting up pycaret", "dataset shape"],
+            "done_kw": ["comparing models", "step: comparing models..."],
+            "description": "Initializes the PyCaret setup, handling normalization, encoding, and train/test splits internally.",
+        },
+        {
+            "label": "Model Comparison",
+            "icon": "⚖️",
+            "keywords": ["comparing models", "including fast/robust models"],
+            "done_kw": ["tuning best model", "step: tuning best model..."],
+            "description": "Trains and evaluates a fast baseline of multiple estimators to find the top candidates.",
+        },
+        {
+            "label": "Hyperparameter Tuning",
+            "icon": "🔧",
+            "keywords": ["tuning best model", "step: tuning best model..."],
+            "done_kw": ["blending top models", "step: blending top models..."],
+            "description": "Applies randomized search to optimize hyperparameters of the best performing model.",
+        },
+        {
+            "label": "Model Blending",
+            "icon": "🌪️",
+            "keywords": ["blending top models", "step: blending top models..."],
+            "done_kw": ["saving model", "pycaret experiment completed"],
+            "description": "Creates an ensemble of the top models to improve generalized performance via voting/averaging.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["saving model to", "pycaret experiment completed", "thread finished"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+    "lale": [
+        {
+            "label": "Pipeline Definition",
+            "icon": "⚙️",
+            "keywords": ["defining lale planned pipeline", "dataset shape"],
+            "done_kw": ["tuning with hyperopt", "step: tuning with hyperopt..."],
+            "description": "Maps a search space over transformers (PCA, Scalers) and estimators (LR, RF, KNN).",
+        },
+        {
+            "label": "Hyperopt Tuning",
+            "icon": "🔧",
+            "keywords": ["tuning with hyperopt", "step: tuning with hyperopt..."],
+            "done_kw": ["fitting lale optimizer", "step: fitting lale optimizer"],
+            "description": "Configures Tree-structured Parzen Estimators (TPE) algorithm for intelligent hyperparameter search.",
+        },
+        {
+            "label": "Fitting Optimizer",
+            "icon": "🕒",
+            "keywords": ["fitting lale optimizer", "step: fitting lale optimizer"],
+            "done_kw": ["saving model locally", "step: saving model locally"],
+            "description": "Executes identical cross-validation folds on generated pipelines within the set budget.",
+        },
+        {
+            "label": "Best Model Extraction",
+            "icon": "🏆",
+            "keywords": ["best pipeline structure:", "best f1 (macro) score"],
+            "done_kw": ["saving model locally", "step: saving model locally"],
+            "description": "Decodes the structure and metrics of the optimized pipeline graph.",
+        },
+        {
+            "label": "MLflow Logging",
+            "icon": "📝",
+            "keywords": ["saving model locally", "lale experiment completed", "thread finished"],
+            "done_kw": ["thread finished"],
+            "description": "Persists model artifacts, parameters, and metrics to MLflow for tracking and versioning.",
+        },
+    ],
+}
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def get_framework_steps(framework_key: str) -> list[dict]:
+    """Return the step definitions for a given framework key."""
+    return _STEPS.get(framework_key.lower(), [])
+
+
+def infer_pipeline_steps(framework_key: str, logs: list[str], status: str) -> list[dict]:
+    """
+    Returns enriched step list with status attached:
+      status = "done" | "active" | "pending"
+
+    On completed/failed/cancelled runs, all matched steps are "done".
+    """
+    steps = get_framework_steps(framework_key)
+    if not steps:
+        return []
+
+    log_blob = " ".join(logs).lower()
+
+    if status == "completed":
+        # Mark all steps done
+        return [{"label": s["label"], "icon": s["icon"], "description": s["description"], "status": "done"} for s in steps]
+
+    if status in ("failed", "cancelled"):
+        # Mark up to the last-seen step as done, rest pending, mark last active as failed
+        last_done_idx = -1
+        for i, step in enumerate(steps):
+            if any(kw in log_blob for kw in step["keywords"]):
+                last_done_idx = i
+
+        result = []
+        for i, step in enumerate(steps):
+            if i < last_done_idx:
+                st_val = "done"
+            elif i == last_done_idx:
+                st_val = "failed" if status == "failed" else "cancelled"
+            else:
+                st_val = "pending"
+            result.append({"label": step["label"], "icon": step["icon"], "description": step["description"], "status": st_val})
+        return result
+
+    # Running or queued: find the active step
+    last_done_idx = -1
+    for i, step in enumerate(steps):
+        done_signals = step.get("done_kw", [])
+        if any(kw in log_blob for kw in done_signals):
+            last_done_idx = i
+
+    # Active = first step after last_done
+    active_idx = min(last_done_idx + 1, len(steps) - 1)
+
+    result = []
+    for i, step in enumerate(steps):
+        if i <= last_done_idx:
+            st_val = "done"
+        elif i == active_idx and status == "running":
+            st_val = "active"
+        else:
+            st_val = "pending"
+        result.append({"label": step["label"], "icon": step["icon"], "description": step["description"], "status": st_val})
+
+    return result
+
+
+def extract_best_tpot_pipeline(logs: list[str]) -> Optional[str]:
+    """Extract the TPOT best pipeline string from logs."""
+    for line in reversed(logs):
+        if "best pipeline:" in line.lower() or "fitted_pipeline_" in line.lower():
+            return line.strip()
+        if "pipeline(" in line.lower():
+            return line.strip()
+    return None
+
+
+def extract_autogluon_leaderboard_text(logs: list[str]) -> Optional[str]:
+    """Extract leaderboard table text from AutoGluon logs."""
+    rows = []
+    capture = False
+    for line in logs:
+        if "model" in line.lower() and "score_val" in line.lower():
+            capture = True
+        if capture:
+            rows.append(line)
+            if len(rows) > 15:
+                break
+    return "\n".join(rows) if rows else None

src/pycaret_utils.py

@@ -0,0 +1,212 @@
+import os
+import logging
+import traceback
+import queue
+import time
+import pandas as pd
+from typing import Dict, Any, Optional
+
+import mlflow
+
+from src.mlflow_utils import safe_set_experiment
+from src.onnx_utils import export_to_onnx
+
+
+def run_pycaret_experiment(
+    train_df: pd.DataFrame,
+    target_col: str,
+    run_name: str,
+    time_limit: Optional[int],
+    log_queue: queue.Queue,
+    stop_event=None,
+    val_df: Optional[pd.DataFrame] = None,
+    task_type: str = "Classification",
+    n_jobs: int = 1,
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    Run PyCaret experiment.
+    Dynamically loads classification, regression, or time_series depending on task_type.
+    """
+    logger = logging.getLogger("pycaret")
+    logger.info(f"Starting PyCaret experiment: {run_name} (Task: {task_type})")
+    logger.info(f"Dataset shape: {train_df.shape}, Target: {target_col}")
+
+    # Dynamic imports based on task_type
+    if task_type == "Regression":
+        from pycaret.regression import setup, compare_models, pull, tune_model, blend_models, save_model
+        sort_metric = "R2"
+        include_models = ["lr", "rf", "et", "lightgbm"]
+    elif task_type == "Time Series Forecasting":
+        from pycaret.time_series import setup, compare_models, pull, tune_model, blend_models, save_model
+        sort_metric = "MASE"
+        include_models = ["naive", "snaive", "arima", "ets"]
+    else:
+        from pycaret.classification import setup, compare_models, pull, tune_model, blend_models, save_model
+        sort_metric = "F1"
+        include_models = ["lr", "nb", "rf", "et", "lightgbm"]
+
+    # Always end any dangling MLflow run to avoid conflicts
+    try:
+        mlflow.end_run()
+    except Exception:
+        pass
+
+    # 1. Prepare MLflow Tracking
+    safe_set_experiment("Multi_AutoML_Project")
+
+    if stop_event and stop_event.is_set():
+        raise StopIteration("Experiment cancelled before setup.")
+
+    try:
+        # 2. PyCaret Setup
+        logger.info("Step: Setting up PyCaret environment...")
+        
+        setup_kwargs = {
+            "data": train_df,
+            "target": target_col,
+            "session_id": 42,
+            "verbose": False,
+            "fold": 3,
+            "log_experiment": False,
+            "system_log": False,
+            "n_jobs": n_jobs
+        }
+        
+        if task_type == "Time Series Forecasting":
+            setup_kwargs["fh"] = kwargs.get("fh", 12)
+            setup_kwargs["seasonal_period"] = kwargs.get("seasonal_period", 12)
+        else:
+            setup_kwargs["test_data"] = val_df
+            setup_kwargs["normalize"] = True
+            setup_kwargs["index"] = False
+            setup_kwargs["feature_selection"] = False
+            setup_kwargs["memory"] = False
+
+        clf_setup = setup(**setup_kwargs)
+
+        if stop_event and stop_event.is_set():
+            raise StopIteration("Experiment cancelled after setup.")
+
+        # 3. Start our own MLflow run AFTER PyCaret setup
+        with mlflow.start_run(run_name=run_name) as run:
+            run_id = run.info.run_id
+            logger.info(f"MLflow Run ID: {run_id}")
+            mlflow.log_param("framework", "pycaret")
+            mlflow.log_param("model_type", "pycaret")
+            mlflow.log_param("task_type", task_type)
+
+            # 4. Model Comparison
+            logger.info("Step: Comparing models...")
+            n_select = 3
+            logger.info(f"Including models: {include_models} (Sorting by {sort_metric})")
+
+            best_models = compare_models(
+                n_select=n_select,
+                sort=sort_metric,
+                verbose=False,
+                include=include_models
+            )
+
+            comparison_df = pull()
+            if not comparison_df.empty:
+                top_model_name = comparison_df.iloc[0]['Model']
+                logger.info(f"Best model found: {top_model_name}")
+
+            if stop_event and stop_event.is_set():
+                raise StopIteration("Experiment cancelled after model comparison.")
+
+            # Ensure best_models is a list
+            if not isinstance(best_models, list):
+                best_models = [best_models]
+
+            best_model = best_models[0]
+
+            # 5. Tuning (Time Series tuning might require different params, keeping generic)
+            logger.info("Step: Tuning best model...")
+            n_iter = 10 if time_limit is None or time_limit >= 300 else 5
+            
+            # search_library="scikit-learn" shouldn't be passed to pycaret.time_series
+            tune_kwargs = {
+                "estimator": best_model,
+                "optimize": sort_metric,
+                "n_iter": n_iter,
+                "verbose": False,
+                "choose_better": True
+            }
+            if task_type != "Time Series Forecasting":
+                tune_kwargs["search_library"] = "scikit-learn"
+                tune_kwargs["search_algorithm"] = "random"
+
+            tuned_model = tune_model(**tune_kwargs)
+
+            if stop_event and stop_event.is_set():
+                raise StopIteration("Experiment cancelled after tuning.")
+
+            # 6. Blending (only if we have multiple models)
+            if len(best_models) > 1:
+                logger.info("Step: Blending top models...")
+                final_model = blend_models(
+                    estimator_list=best_models,
+                    optimize=sort_metric,
+                    verbose=False
+                )
+            else:
+                final_model = tuned_model
+                logger.info("Step: Skipping blend (only one model selected).")
+
+            # 7. Save model
+            model_dir = "models"
+            os.makedirs(model_dir, exist_ok=True)
+            model_path_base = os.path.join(model_dir, f"{run_name}_pycaret_model")
+            logger.info(f"Saving model to {model_path_base}.pkl...")
+            save_model(final_model, model_path_base)
+
+            # 8. Log metrics to our MLflow run
+            try:
+                final_metrics = pull()
+                if not final_metrics.empty:
+                    row = final_metrics.iloc[0]
+                    for k, v in row.items():
+                        if isinstance(v, (int, float)):
+                            mlflow.log_metric(k.lower().replace(" ", "_"), float(v))
+            except Exception as me:
+                logger.warning(f"Could not pull metrics: {me}")
+
+            # Log model artifact
+            model_pkl = f"{model_path_base}.pkl"
+            if os.path.exists(model_pkl):
+                mlflow.log_artifact(model_pkl, artifact_path="model")
+
+            # ONNX Export
+            try:
+                onnx_path = os.path.join(model_dir, f"{run_name}_pycaret.onnx")
+                # PyCaret 'final_model' is a scikit-learn pipeline
+                export_to_onnx(final_model, "pycaret", target_col, onnx_path, input_sample=train_df[:1])
+                mlflow.log_artifact(onnx_path, artifact_path="model")
+            except Exception as e:
+                logger.warning(f"Failed to export PyCaret model to ONNX: {e}")
+
+            logger.info("PyCaret experiment completed successfully.")
+            return {
+                "success": True,
+                "predictor": final_model,
+                "run_id": run_id,
+                "type": "pycaret",
+                "model_path": model_pkl
+            }
+
+    except StopIteration as si:
+        logger.warning(f"Cancelled: {si}")
+        raise
+    except Exception as e:
+        logger.error(f"PyCaret Error: {e}")
+        logger.error(traceback.format_exc())
+        raise e
+    finally:
+        # Always clean up any dangling run
+        try:
+            mlflow.end_run()
+        except Exception:
+            pass
+

src/tpot_utils.py

@@ -165,7 +165,8 @@ def train_tpot_model(df, target_column, run_name,
                      generations=5, population_size=20, cv=5, 
                      scoring=None, max_time_mins=30, max_eval_time_mins=5, random_state=42, 
                      verbosity=2, n_jobs=-1, config_dict='TPOT sparse',
-                     tfidf_max_features=500, tfidf_ngram_range=(1, 2)):
+                     tfidf_max_features=500, tfidf_ngram_range=(1, 2),
+                     stop_event=None):
     """
     Train TPOT model with MLflow tracking
     """
@@ -232,10 +233,13 @@ def train_tpot_model(df, target_column, run_name,
                 scoring = 'neg_mean_squared_error'
         
         # Ensure there are no loose active runs that could cause errors on start
-        while mlflow.active_run():
-            mlflow.end_run()
+        try:
+            while mlflow.active_run():
+                mlflow.end_run()
+        except:
+            pass
             
-        with mlflow.start_run(run_name=run_name) as run:
+        with mlflow.start_run(run_name=run_name, nested=True) as run:
             logger.info(f"Starting TPOT training for run: {run_name}")
             
             # Choose TPOT class based on problem type
@@ -432,6 +436,19 @@ def train_tpot_model(df, target_column, run_name,
             # Log the fitted pipeline
             mlflow.sklearn.log_model(final_pipeline, "model", registered_model_name=f"TPOT_{run_name}")
             
+            # Generate and log consumption code sample
+            try:
+                from src.code_gen_utils import generate_consumption_code
+                code_sample = generate_consumption_code("tpot", run.info.run_id, target_column)
+                code_path = "consumption_sample.py"
+                with open(code_path, "w") as f:
+                    f.write(code_sample)
+                mlflow.log_artifact(code_path)
+                if os.path.exists(code_path):
+                    os.remove(code_path)
+            except Exception as e:
+                logger.warning(f"Failed to generate consumption code: {e}")
+
             logger.info("TPOT model successfully registered in MLflow")
             
             return tpot, final_pipeline, run.info.run_id, model_info

src/training_worker.py

@@ -0,0 +1,218 @@
+"""
+training_worker: thread entry point for every AutoML run.
+Captures stdout/stderr, feeds log_queue, puts result into result_queue,
+and respects the stop_event for graceful cancellation.
+
+Log isolation strategy (definitive):
+  - We attach a _QueueLogHandler to each relevant named library logger.
+  - Each handler has a _ThreadFilter that only accepts log records whose
+    record.thread matches the experiment thread's ID.
+  - This means messages from Thread A never land in Thread B's queue,
+    even though they share the same named logger objects.
+  - propagate is set to False to prevent double-delivery via the root logger.
+  - All are restored in the finally block.
+
+Stdout/Stderr isolation:
+  - redirect_stdout/redirect_stderr are process-global (they overwrite sys.stdout).
+  - We use a _ThreadAwareIO wrapper instead: it checks threading.current_thread()
+    on every write() call, so writes only reach the owning thread's queue.
+"""
+import io
+import sys
+import logging
+import threading
+import traceback
+
+from src.experiment_manager import ExperimentEntry
+
+_LIB_LOGGERS = [
+    'flaml', 'autogluon', 'mlflow', 'h2o', 'tpot',
+    'pycaret', 'lale', 'hyperopt', 'lightgbm', 'xgboost', 'catboost'
+]
+
+
+# ---------------------------------------------------------------------------
+# Thread-aware stdout/stderr router (installed once, process-wide)
+# ---------------------------------------------------------------------------
+class _ThreadAwareIO(io.TextIOBase):
+    """
+    Drop-in replacement for sys.stdout / sys.stderr that routes each write()
+    to the queue registered for the current thread, or falls back to the
+    original stream.
+    """
+    def __init__(self, original_stream):
+        super().__init__()
+        self._original = original_stream
+        self._lock = threading.Lock()
+        self._thread_queues: dict[int, "queue.Queue"] = {}
+
+    def register(self, thread_id: int, q):
+        with self._lock:
+            self._thread_queues[thread_id] = q
+
+    def unregister(self, thread_id: int):
+        with self._lock:
+            self._thread_queues.pop(thread_id, None)
+
+    def write(self, s: str) -> int:
+        if not isinstance(s, str):
+            try:
+                s = str(s)
+            except Exception:
+                return 0
+        tid = threading.current_thread().ident
+        with self._lock:
+            q = self._thread_queues.get(tid)
+        if q is not None:
+            if s.strip():
+                # Filter out progress bar characters that fail on Windows cp1252
+                # \u2588 is the full block character
+                safe_s = s.replace('\u2588', '#').replace('\u258c', '|').replace('\u2584', '-')
+                q.put(safe_s.strip())
+        else:
+            # Fall back to original stream for threads not registered
+            try:
+                self._original.write(s)
+            except Exception:
+                pass
+        return len(s)
+
+    def flush(self):
+        try:
+            self._original.flush()
+        except Exception:
+            pass
+
+    @property
+    def encoding(self):
+        return getattr(self._original, 'encoding', 'utf-8') or 'utf-8'
+
+    @property
+    def errors(self):
+        return getattr(self._original, 'errors', 'replace')
+
+
+# Install thread-aware routers once for the entire process
+_stdout_router = _ThreadAwareIO(sys.__stdout__)
+_stderr_router = _ThreadAwareIO(sys.__stderr__)
+sys.stdout = _stdout_router
+sys.stderr = _stderr_router
+
+
+# ---------------------------------------------------------------------------
+# Per-thread log handler with thread filter
+# ---------------------------------------------------------------------------
+class _ThreadFilter(logging.Filter):
+    """Only accepts log records emitted by a specific OS thread."""
+    def __init__(self, thread_id: int):
+        super().__init__()
+        self._thread_id = thread_id
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        return record.thread == self._thread_id
+
+
+class _QueueLogHandler(logging.Handler):
+    def __init__(self, log_queue):
+        super().__init__()
+        self.log_queue = log_queue
+
+    def emit(self, record: logging.LogRecord):
+        try:
+            msg = self.format(record)
+            self.log_queue.put(msg)
+        except Exception:
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Worker entry point
+# ---------------------------------------------------------------------------
+def run_training_worker(entry: ExperimentEntry, train_fn, kwargs: dict):
+    """
+    Thread target. Runs train_fn(**kwargs, stop_event=entry.stop_event),
+    keeps the entry status updated, and puts the final result dict in
+    result_queue.
+    """
+    thread_id = threading.current_thread().ident
+
+    # --- Thread-aware stdout/stderr capture ---
+    _stdout_router.register(thread_id, entry.log_queue)
+    _stderr_router.register(thread_id, entry.log_queue)
+
+    # --- Per-thread logging handler (with thread filter) ---
+    handler = _QueueLogHandler(entry.log_queue)
+    handler.setFormatter(logging.Formatter('%(message)s'))
+    handler.addFilter(_ThreadFilter(thread_id))
+
+    saved_propagate: dict[str, bool] = {}
+    for lib in _LIB_LOGGERS:
+        lib_logger = logging.getLogger(lib)
+        saved_propagate[lib] = lib_logger.propagate
+        lib_logger.propagate = False  # prevents root from seeing AND double-deliver
+        lib_logger.addHandler(handler)
+        if lib_logger.level == logging.NOTSET or lib_logger.level > logging.INFO:
+            lib_logger.setLevel(logging.INFO)
+
+    entry.status = "running"
+    entry.log_queue.put(f"[Worker] Starting training: {entry.metadata.get('run_name', entry.key)}")
+
+    try:
+        # Inject stop_event and telemetry_queue into kwargs if the function accepts it
+        try:
+            import inspect
+            sig = inspect.signature(train_fn)
+            if 'stop_event' in sig.parameters:
+                kwargs['stop_event'] = entry.stop_event
+            if 'telemetry_queue' in sig.parameters:
+                kwargs['telemetry_queue'] = entry.telemetry_queue
+        except Exception:
+            pass
+
+        result = train_fn(**kwargs)
+
+        # Normalise result into a standard dict
+        if isinstance(result, tuple):
+            if len(result) == 2:
+                predictor, run_id = result
+                entry.result_queue.put({
+                    "success": True, "predictor": predictor, "run_id": run_id,
+                    "type": entry.metadata.get("framework_key", "unknown")
+                })
+            elif len(result) == 4:
+                tpot, pipeline, run_id, info = result
+                entry.result_queue.put({
+                    "success": True, "predictor": pipeline, "run_id": run_id, "info": info, "type": "tpot"
+                })
+            else:
+                entry.result_queue.put({
+                    "success": True, "predictor": result[0], "run_id": result[-1],
+                    "type": entry.metadata.get("framework_key", "unknown")
+                })
+        elif isinstance(result, dict):
+            entry.result_queue.put(result)
+        else:
+            entry.result_queue.put({
+                "success": True, "predictor": result, "run_id": None,
+                "type": entry.metadata.get("framework_key", "unknown")
+            })
+
+    except StopIteration:
+        entry.log_queue.put("[Worker] Training cancelled by user request.")
+        entry.result_queue.put({"success": False, "cancelled": True, "error": "Cancelled by user"})
+    except Exception as e:
+        err_tb = traceback.format_exc()
+        entry.log_queue.put(f"[Worker] CRITICAL ERROR: {e}\n{err_tb}")
+        entry.result_queue.put({"success": False, "error": str(e), "traceback": err_tb})
+    finally:
+        # Restore all lib loggers
+        for lib in _LIB_LOGGERS:
+            lib_logger = logging.getLogger(lib)
+            lib_logger.removeHandler(handler)
+            lib_logger.propagate = saved_propagate.get(lib, True)
+
+        # Unregister stdout/stderr routing for this thread
+        _stdout_router.unregister(thread_id)
+        _stderr_router.unregister(thread_id)
+
+        entry.log_queue.put("[Worker] Thread finished.")

src/xai_utils.py

@@ -0,0 +1,216 @@
+import os
+import io
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import warnings
+
+def generate_shap_explanation(model, X_train: pd.DataFrame, X_valid: pd.DataFrame = None, 
+                              max_background_samples=100, task_type="Classification"):
+    """
+    Generates SHAP Global Feature Importance plot for Tabular data.
+    """
+    try:
+        import shap
+    except ImportError:
+        warnings.warn("SHAP library not installed. Cannot generate explanations.")
+        return None
+
+    plt.switch_backend('Agg') # Ensure thread-safe rendering without GUI
+
+    # 1. Determine background dataset (handle large data gracefully)
+    bg_data = X_train
+    if len(bg_data) > max_background_samples:
+        bg_data = bg_data.sample(n=max_background_samples, random_state=42)
+        
+    evaluate_data = X_valid if X_valid is not None else bg_data
+    if len(evaluate_data) > max_background_samples:
+        evaluate_data = evaluate_data.sample(n=max_background_samples, random_state=42)
+
+    # Convert non-numeric for generic shap handling if required by models
+    # Depending on framework, categorical columns might need Ordinal/OneHot. 
+    # For robust black-box generic explainer:
+    
+    explainer = None
+    shap_values = None
+    
+    # 2. Heuristics to pick the right explainer
+    model_type = str(type(model)).lower()
+    
+    try:
+        if 'lgbm' in model_type or 'xgb' in model_type or 'catboost' in model_type or 'ensemble' in model_type:
+            # TreeExplainer is fast for tree-based models and forests
+            try:
+                explainer = shap.TreeExplainer(model)
+                shap_values = explainer.shap_values(evaluate_data)
+            except Exception:
+                pass # Fallback to generic
+
+        if explainer is None:
+            # For complex pipelines (like sklearn pipelines, PyCaret, generic wrappers)
+            # Use KernelExplainer as a Black-Box proxy (requires a predict function)
+            
+            predict_fn = None
+            if hasattr(model, "predict_proba") and "classification" in task_type.lower():
+                predict_fn = lambda x: model.predict_proba(x)
+            elif hasattr(model, "predict"):
+                predict_fn = lambda x: model.predict(x)
+            else:
+                return None # Can't explain
+                
+            # KernelExplainer can be slow, hence the small bg_data
+            explainer = shap.KernelExplainer(predict_fn, bg_data)
+            shap_values = explainer.shap_values(evaluate_data)
+            
+    except Exception as e:
+        warnings.warn(f"SHAP generation failed: {e}")
+        return None
+
+    # 3. Generate the Plot
+    fig = plt.figure(figsize=(10, 6))
+    
+    try:
+        # For multi-class, shap_values is a list. For regression/binary, it's an array.
+        if isinstance(shap_values, list):
+            # Take the shap values for the first class/positive class for overview
+            shap.summary_plot(shap_values[1] if len(shap_values)>1 else shap_values[0], evaluate_data, show=False)
+        else:
+            shap.summary_plot(shap_values, evaluate_data, show=False)
+            
+        plt.tight_layout()
+        return fig
+    except Exception as e:
+        warnings.warn(f"SHAP plot rendering failed: {e}")
+        plt.close(fig)
+        return None
+
+def generate_cv_saliency_map(model, image_path: str, target_size=(224, 224), step=15, window_size=30):
+    """
+    Universal Occlusion Saliency Map for Black-Box CV Models (AutoGluon/AutoKeras).
+    Instead of relying on internal hooks (which heavily abstracted AutoML layers hide),
+    we slide a black box ('occlusion') across the image and measure the confidence drop.
+    The regions that drop the confidence the most are the most salient (important) for the prediction.
+    """
+    try:
+        from PIL import Image
+        import cv2
+    except ImportError:
+        warnings.warn("Missing CV libraries (Pillow/OpenCV) for Saliency representation.")
+        return None
+
+    try:
+        # 1. Load Original Image
+        original_img = Image.open(image_path).convert('RGB')
+        img_w, img_h = original_img.size
+        
+        # Determine the baseline prediction to see what class we are explaining
+        # Since this is a generic AutoML predictor UI, we assume `model.predict_proba` gives a df or dict
+        df_single = pd.DataFrame([{"image": image_path}])
+        
+        # Get base probabilities. 
+        # Note: Depending on AutoGluon/AutoKeras formatting, the predict_proba method might vary.
+        if hasattr(model, 'predict_proba'):
+            base_probs = model.predict_proba(df_single)
+            if isinstance(base_probs, pd.DataFrame):
+                # Assuming top class 
+                top_class = base_probs.iloc[0].idxmax()
+                base_score = base_probs.iloc[0][top_class]
+            else:
+                top_class = np.argmax(base_probs[0])
+                base_score = base_probs[0][top_class]
+        else:
+            warnings.warn("Model does not support predict_proba, Saliency Map cannot track confidence drops.")
+            return None
+
+        # 2. Build Saliency Map Array
+        saliency_map = np.zeros((img_h, img_w))
+        heatmap_counts = np.zeros((img_h, img_w))
+
+        # We will create occluded images, save them temporarily, and batch-predict to find drops
+        # For performance, we downsize the grid if the image is huge
+        grid_step = step
+        w_size = window_size
+        
+        # To avoid predicting 1000s of images, let's limit the grid
+        if (img_h / step) * (img_w / step) > 200:
+            grid_step = max(int(img_h/10), 10)
+            w_size = int(grid_step * 1.5)
+
+        occluded_paths = []
+        coords = []
+        
+        tmp_dir = os.path.join("data_lake", "tmp_occlusion")
+        os.makedirs(tmp_dir, exist_ok=True)
+        img_arr_orig = np.array(original_img)
+
+        # Generate Occluded Copies
+        for y in range(0, img_h, grid_step):
+            for x in range(0, img_w, grid_step):
+                img_copy = img_arr_orig.copy()
+                
+                # Apply black box
+                y1, y2 = max(0, y - w_size // 2), min(img_h, y + w_size // 2)
+                x1, x2 = max(0, x - w_size // 2), min(img_w, x + w_size // 2)
+                img_copy[y1:y2, x1:x2] = 0 # Occlude
+                
+                t_path = os.path.join(tmp_dir, f"occ_{y}_{x}.jpg")
+                Image.fromarray(img_copy).save(t_path)
+                occluded_paths.append(t_path)
+                coords.append((y1, y2, x1, x2))
+
+        # Predict all simultaneously
+        df_batch = pd.DataFrame({"image": occluded_paths})
+        
+        try:
+            batch_probs = model.predict_proba(df_batch)
+        except Exception:
+            warnings.warn("Batch probability prediction failed for occlusion map.")
+            return None
+            
+        # Parse scores based on framework signature
+        if isinstance(batch_probs, pd.DataFrame):
+            scores = batch_probs[top_class].values
+        else:
+            scores = batch_probs[:, top_class] if len(batch_probs.shape) > 1 else batch_probs
+
+        # 3. Calculate importance based on score drops
+        for idx, (y1, y2, x1, x2) in enumerate(coords):
+            drop = base_score - scores[idx]
+            # If the score dropped, this region was important
+            importance = max(0, drop)
+            saliency_map[y1:y2, x1:x2] += importance
+            heatmap_counts[y1:y2, x1:x2] += 1
+
+        # Average overlaps
+        heatmap_counts[heatmap_counts == 0] = 1
+        saliency_avg = saliency_map / heatmap_counts
+
+        # Normalize 0-255
+        if np.max(saliency_avg) > 0:
+            saliency_avg = (saliency_avg / np.max(saliency_avg)) * 255
+        saliency_avg = np.uint8(saliency_avg)
+
+        # 4. Generate visual overlay
+        colormap = cv2.applyColorMap(saliency_avg, cv2.COLORMAP_JET)
+        
+        orig_cv = cv2.cvtColor(np.array(original_img), cv2.COLORRGB_BGR) # To match cv2
+        final_overlay = cv2.addWeighted(orig_cv, 0.6, colormap, 0.4, 0)
+        final_rgb = cv2.cvtColor(final_overlay, cv2.COLORBGR_RGB)
+        
+        # Cleanup
+        for p in occluded_paths:
+            try: os.remove(p)
+            except: pass
+
+        fig = plt.figure(figsize=(8, 8))
+        plt.imshow(final_rgb)
+        plt.title(f"XAI Occlusion Heatmap (Target: {top_class})")
+        plt.axis('off')
+        plt.tight_layout()
+        
+        return fig
+
+    except Exception as e:
+        warnings.warn(f"CV XAI generation failed: {e}")
+        return None