ins-pricing 0.1.11__py3-none-any.whl → 0.2.0__py3-none-any.whl

ins_pricing_gemini/modelling/explain/permutation.py

@@ -1,155 +0,0 @@
-from __future__ import annotations
-
-from typing import Callable, Optional, Sequence
-
-import numpy as np
-import pandas as pd
-
-from .metrics import resolve_metric
-
-
-def _prepare_data(X, y, sample_weight, max_rows, rng):
-    y_arr = np.asarray(y)
-    if y_arr.ndim != 1:
-        y_arr = y_arr.reshape(-1)
-
-    w_arr = None
-    if sample_weight is not None:
-        w_arr = np.asarray(sample_weight).reshape(-1)
-        if w_arr.shape[0] != y_arr.shape[0]:
-            raise ValueError("sample_weight length must match y.")
-
-    if isinstance(X, pd.DataFrame):
-        X_data = X
-        if len(X_data) != len(y_arr):
-            raise ValueError("X and y must have the same length.")
-        if max_rows and len(X_data) > max_rows:
-            idx = rng.choice(len(X_data), size=int(max_rows), replace=False)
-            X_data = X_data.iloc[idx].copy()
-            y_arr = y_arr[idx]
-            if w_arr is not None:
-                w_arr = w_arr[idx]
-        return X_data, y_arr, w_arr
-
-    X_np = np.asarray(X)
-    if X_np.ndim != 2:
-        raise ValueError("X must be 2d when not a DataFrame.")
-    if X_np.shape[0] != y_arr.shape[0]:
-        raise ValueError("X and y must have the same length.")
-    if max_rows and X_np.shape[0] > max_rows:
-        idx = rng.choice(X_np.shape[0], size=int(max_rows), replace=False)
-        X_np = X_np[idx]
-        y_arr = y_arr[idx]
-        if w_arr is not None:
-            w_arr = w_arr[idx]
-    return X_np, y_arr, w_arr
-
-
-def permutation_importance(
-    predict_fn: Callable,
-    X,
-    y,
-    *,
-    sample_weight=None,
-    metric: str | Callable = "auto",
-    task_type: Optional[str] = None,
-    higher_is_better: Optional[bool] = None,
-    n_repeats: int = 5,
-    random_state: Optional[int] = None,
-    max_rows: Optional[int] = 5000,
-    features: Optional[Sequence[str]] = None,
-    return_scores: bool = False,
-    safe_copy: bool = False,
-) -> pd.DataFrame:
-    """Permutation importance on tabular data.
-
-    predict_fn should accept the same type as X (DataFrame or ndarray).
-    Set safe_copy=True if predict_fn mutates its input.
-    """
-    rng = np.random.default_rng(random_state)
-    n_repeats = max(1, int(n_repeats))
-
-    X_data, y_arr, w_arr = _prepare_data(X, y, sample_weight, max_rows, rng)
-    metric_fn, higher_is_better, metric_name = resolve_metric(
-        metric, task_type=task_type, higher_is_better=higher_is_better
-    )
-
-    baseline_pred = predict_fn(X_data)
-    baseline_score = metric_fn(y_arr, baseline_pred, w_arr)
-
-    if isinstance(X_data, pd.DataFrame):
-        feature_names = list(X_data.columns)
-        if features is not None:
-            feature_names = [f for f in features if f in X_data.columns]
-        X_perm = X_data.copy()
-        results = []
-        for feat in feature_names:
-            orig_series = X_perm[feat].copy()
-            orig_values = orig_series.to_numpy(copy=True)
-            scores = []
-            for _ in range(n_repeats):
-                X_perm[feat] = rng.permutation(orig_values)
-                pred_input = X_perm.copy() if safe_copy else X_perm
-                pred = predict_fn(pred_input)
-                score = metric_fn(y_arr, pred, w_arr)
-                scores.append(float(score))
-            X_perm[feat] = orig_series
-
-            scores_arr = np.asarray(scores, dtype=float)
-            if higher_is_better:
-                delta = baseline_score - scores_arr
-            else:
-                delta = scores_arr - baseline_score
-            entry = {
-                "feature": feat,
-                "importance_mean": float(np.mean(delta)),
-                "importance_std": float(np.std(delta)),
-                "baseline_score": float(baseline_score),
-                "permutation_score_mean": float(np.mean(scores_arr)),
-                "metric": metric_name,
-            }
-            if return_scores:
-                entry["permutation_scores"] = scores
-            results.append(entry)
-    else:
-        if features is not None:
-            if len(features) != X_data.shape[1]:
-                raise ValueError("features length must match X columns for ndarray input.")
-            feature_names = list(features)
-        else:
-            feature_names = [f"x{i}" for i in range(X_data.shape[1])]
-
-        X_base = np.asarray(X_data)
-        X_perm = X_base.copy()
-        results = []
-        for idx, feat in enumerate(feature_names):
-            orig_col = X_base[:, idx].copy()
-            scores = []
-            for _ in range(n_repeats):
-                X_perm[:, idx] = rng.permutation(orig_col)
-                pred_input = X_perm.copy() if safe_copy else X_perm
-                pred = predict_fn(pred_input)
-                score = metric_fn(y_arr, pred, w_arr)
-                scores.append(float(score))
-            X_perm[:, idx] = orig_col
-
-            scores_arr = np.asarray(scores, dtype=float)
-            if higher_is_better:
-                delta = baseline_score - scores_arr
-            else:
-                delta = scores_arr - baseline_score
-            entry = {
-                "feature": feat,
-                "importance_mean": float(np.mean(delta)),
-                "importance_std": float(np.std(delta)),
-                "baseline_score": float(baseline_score),
-                "permutation_score_mean": float(np.mean(scores_arr)),
-                "metric": metric_name,
-            }
-            if return_scores:
-                entry["permutation_scores"] = scores
-            results.append(entry)
-
-    df = pd.DataFrame(results)
-    df = df.sort_values(by="importance_mean", ascending=False).reset_index(drop=True)
-    return df

ins_pricing_gemini/modelling/explain/shap_utils.py

@@ -1,146 +0,0 @@
-from __future__ import annotations
-
-from typing import Callable, Optional
-
-import numpy as np
-import pandas as pd
-
-
-def _require_shap():
-    try:
-        import shap  # type: ignore
-    except Exception as exc:  # pragma: no cover - optional dependency
-        raise ImportError("SHAP is required. Install with `pip install shap`.") from exc
-    return shap
-
-
-def compute_shap_core(
-    ctx,
-    model_key: str,
-    n_background: int,
-    n_samples: int,
-    on_train: bool,
-    X_df: pd.DataFrame,
-    prep_fn: Callable[[pd.DataFrame], np.ndarray],
-    predict_fn: Callable[[np.ndarray], np.ndarray],
-    cleanup_fn: Optional[Callable[[], None]] = None,
-) -> dict:
-    """Shared SHAP pipeline using KernelExplainer with lazy import."""
-    _ = on_train
-    if model_key not in ctx.trainers or ctx.trainers[model_key].model is None:
-        raise RuntimeError(f"Model {model_key} not trained.")
-    if cleanup_fn:
-        cleanup_fn()
-    shap = _require_shap()
-    bg_df = ctx._sample_rows(X_df, n_background)
-    bg_mat = prep_fn(bg_df)
-    explainer = shap.KernelExplainer(predict_fn, bg_mat)
-    ex_df = ctx._sample_rows(X_df, n_samples)
-    ex_mat = prep_fn(ex_df)
-    nsample_eff = ctx._shap_nsamples(ex_mat)
-    shap_values = explainer.shap_values(ex_mat, nsamples=nsample_eff)
-    bg_pred = predict_fn(bg_mat)
-    base_value = float(np.asarray(bg_pred).mean())
-
-    return {
-        "explainer": explainer,
-        "X_explain": ex_df,
-        "shap_values": shap_values,
-        "base_value": base_value,
-    }
-
-
-def compute_shap_glm(ctx, n_background: int = 500, n_samples: int = 200, on_train: bool = True):
-    data = ctx.train_oht_scl_data if on_train else ctx.test_oht_scl_data
-    design_all = ctx._build_glm_design(data)
-    design_cols = list(design_all.columns)
-
-    def predict_wrapper(x_np):
-        x_df = pd.DataFrame(x_np, columns=design_cols)
-        y_pred = ctx.glm_best.predict(x_df)
-        return np.asarray(y_pred, dtype=np.float64).reshape(-1)
-
-    return compute_shap_core(
-        ctx,
-        "glm",
-        n_background,
-        n_samples,
-        on_train,
-        X_df=design_all,
-        prep_fn=lambda df: df.to_numpy(dtype=np.float64),
-        predict_fn=predict_wrapper,
-    )
-
-
-def compute_shap_xgb(ctx, n_background: int = 500, n_samples: int = 200, on_train: bool = True):
-    data = ctx.train_data if on_train else ctx.test_data
-    X_raw = data[ctx.factor_nmes]
-
-    def predict_wrapper(x_mat):
-        df_input = ctx._decode_ft_shap_matrix_to_df(x_mat)
-        return ctx.xgb_best.predict(df_input)
-
-    return compute_shap_core(
-        ctx,
-        "xgb",
-        n_background,
-        n_samples,
-        on_train,
-        X_df=X_raw,
-        prep_fn=lambda df: ctx._build_ft_shap_matrix(df).astype(np.float64),
-        predict_fn=predict_wrapper,
-    )
-
-
-def compute_shap_resn(ctx, n_background: int = 500, n_samples: int = 200, on_train: bool = True):
-    data = ctx.train_oht_scl_data if on_train else ctx.test_oht_scl_data
-    X = data[ctx.var_nmes]
-
-    def cleanup():
-        import torch
-
-        ctx.resn_best.device = torch.device("cpu")
-        ctx.resn_best.resnet.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    return compute_shap_core(
-        ctx,
-        "resn",
-        n_background,
-        n_samples,
-        on_train,
-        X_df=X,
-        prep_fn=lambda df: df.to_numpy(dtype=np.float64),
-        predict_fn=lambda x: ctx._resn_predict_wrapper(x),
-        cleanup_fn=cleanup,
-    )
-
-
-def compute_shap_ft(ctx, n_background: int = 500, n_samples: int = 200, on_train: bool = True):
-    if str(ctx.config.ft_role) != "model":
-        raise RuntimeError(
-            "FT is configured as embedding-only (ft_role != 'model'); FT SHAP is disabled."
-        )
-    data = ctx.train_data if on_train else ctx.test_data
-    X_raw = data[ctx.factor_nmes]
-
-    def cleanup():
-        import torch
-
-        ctx.ft_best.device = torch.device("cpu")
-        ctx.ft_best.ft.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-
-    return compute_shap_core(
-        ctx,
-        "ft",
-        n_background,
-        n_samples,
-        on_train,
-        X_df=X_raw,
-        prep_fn=lambda df: ctx._build_ft_shap_matrix(df).astype(np.float64),
-        predict_fn=ctx._ft_shap_predict_wrapper,
-        cleanup_fn=cleanup,
-    )

ins_pricing_gemini/modelling/features.py

@@ -1,215 +0,0 @@
-from __future__ import annotations
-
-import pandas as pd
-import numpy as np
-from sklearn.preprocessing import StandardScaler
-from typing import Optional, Dict, Any, TYPE_CHECKING
-import sys
-
-if TYPE_CHECKING:
-    from .core import BayesOptModel
-
-def _add_region_effect(model: "BayesOptModel") -> None:
-    """Partial pooling over province/city to create a smoothed region_effect feature."""
-    prov_col = model.config.region_province_col
-    city_col = model.config.region_city_col
-    if not prov_col or not city_col:
-        return
-    for col in [prov_col, city_col]:
-        if col not in model.train_data.columns:
-            print(f"[RegionEffect] Missing column {col}; skipped.")
-            return
-
-    def safe_mean(df: pd.DataFrame) -> float:
-        w = df[model.weight_nme]
-        y = df[model.resp_nme]
-        # EPS is imported from utils usually, but we can define local or import
-        denom = max(float(w.sum()), 1e-6)
-        return float((y * w).sum() / denom)
-
-    global_mean = safe_mean(model.train_data)
-    alpha = max(float(model.config.region_effect_alpha), 0.0)
-
-    w_all = model.train_data[model.weight_nme]
-    y_all = model.train_data[model.resp_nme]
-    yw_all = y_all * w_all
-
-    prov_sumw = w_all.groupby(model.train_data[prov_col]).sum()
-    prov_sumyw = yw_all.groupby(model.train_data[prov_col]).sum()
-    prov_mean = (prov_sumyw / prov_sumw.clip(lower=1e-6)).astype(float)
-    prov_mean = prov_mean.fillna(global_mean)
-
-    city_sumw = model.train_data.groupby([prov_col, city_col])[
-        model.weight_nme].sum()
-    city_sumyw = yw_all.groupby(
-        [model.train_data[prov_col], model.train_data[city_col]]).sum()
-    city_df = pd.DataFrame({
-        "sum_w": city_sumw,
-        "sum_yw": city_sumyw,
-    })
-    city_df["prior"] = city_df.index.get_level_values(0).map(
-        prov_mean).fillna(global_mean)
-    city_df["effect"] = (
-        city_df["sum_yw"] + alpha * city_df["prior"]
-    ) / (city_df["sum_w"] + alpha).clip(lower=1e-6)
-    city_effect = city_df["effect"]
-
-    def lookup_effect(df: pd.DataFrame) -> pd.Series:
-        idx = pd.MultiIndex.from_frame(df[[prov_col, city_col]])
-        effects = city_effect.reindex(idx).to_numpy(dtype=np.float64)
-        prov_fallback = df[prov_col].map(
-            prov_mean).fillna(global_mean).to_numpy(dtype=np.float64)
-        effects = np.where(np.isfinite(effects), effects, prov_fallback)
-        effects = np.where(np.isfinite(effects), effects, global_mean)
-        return pd.Series(effects, index=df.index, dtype=np.float32)
-
-    re_train = lookup_effect(model.train_data)
-    re_test = lookup_effect(model.test_data)
-
-    col_name = "region_effect"
-    model.train_data[col_name] = re_train
-    model.test_data[col_name] = re_test
-
-    # Sync into one-hot and scaled variants.
-    for df in [model.train_oht_data, model.test_oht_data]:
-        if df is not None:
-            df[col_name] = re_train if df is model.train_oht_data else re_test
-
-    # Standardize region_effect and propagate.
-    scaler = StandardScaler()
-    re_train_s = scaler.fit_transform(
-        re_train.values.reshape(-1, 1)).astype(np.float32).reshape(-1)
-    re_test_s = scaler.transform(
-        re_test.values.reshape(-1, 1)).astype(np.float32).reshape(-1)
-    for df in [model.train_oht_scl_data, model.test_oht_scl_data]:
-        if df is not None:
-            df[col_name] = re_train_s if df is model.train_oht_scl_data else re_test_s
-
-    # Update feature lists.
-    if col_name not in model.factor_nmes:
-        model.factor_nmes.append(col_name)
-    if col_name not in model.num_features:
-        model.num_features.append(col_name)
-    if model.train_oht_scl_data is not None:
-        excluded = {model.weight_nme, model.resp_nme}
-        model.var_nmes = [
-            col for col in model.train_oht_scl_data.columns if col not in excluded
-        ]
-
-def _build_geo_tokens(model: "BayesOptModel", params_override: Optional[Dict[str, Any]] = None):
-    """Internal builder; allows trial overrides and returns None on failure."""
-    from .models import GraphNeuralNetSklearn # lazy import to avoid circle if models imports core
-    
-    geo_cols = list(model.config.geo_feature_nmes or [])
-    if not geo_cols:
-        return None
-
-    available = [c for c in geo_cols if c in model.train_data.columns]
-    if not available:
-        return None
-
-    # Preprocess text/numeric: fill numeric with median, label-encode text, map unknowns.
-    proc_train = {}
-    proc_test = {}
-    for col in available:
-        s_train = model.train_data[col]
-        s_test = model.test_data[col]
-        if pd.api.types.is_numeric_dtype(s_train):
-            tr = pd.to_numeric(s_train, errors="coerce")
-            te = pd.to_numeric(s_test, errors="coerce")
-            med = np.nanmedian(tr)
-            proc_train[col] = np.nan_to_num(tr, nan=med).astype(np.float32)
-            proc_test[col] = np.nan_to_num(te, nan=med).astype(np.float32)
-        else:
-            cats = pd.Categorical(s_train.astype(str))
-            tr_codes = cats.codes.astype(np.float32, copy=True)
-            tr_codes[tr_codes < 0] = len(cats.categories)
-            te_cats = pd.Categorical(
-                s_test.astype(str), categories=cats.categories)
-            te_codes = te_cats.codes.astype(np.float32, copy=True)
-            te_codes[te_codes < 0] = len(cats.categories)
-            proc_train[col] = tr_codes
-            proc_test[col] = te_codes
-
-    train_geo_raw = pd.DataFrame(proc_train, index=model.train_data.index)
-    test_geo_raw = pd.DataFrame(proc_test, index=model.test_data.index)
-
-    scaler = StandardScaler()
-    train_geo = pd.DataFrame(
-        scaler.fit_transform(train_geo_raw),
-        columns=available,
-        index=model.train_data.index
-    )
-    test_geo = pd.DataFrame(
-        scaler.transform(test_geo_raw),
-        columns=available,
-        index=model.test_data.index
-    )
-
-    tw_power = model.default_tweedie_power()
-
-    cfg = params_override or {}
-    try:
-        geo_gnn = GraphNeuralNetSklearn(
-            model_nme=f"{model.model_nme}_geo",
-            input_dim=len(available),
-            hidden_dim=cfg.get("geo_token_hidden_dim",
-                               model.config.geo_token_hidden_dim),
-            num_layers=cfg.get("geo_token_layers",
-                               model.config.geo_token_layers),
-            k_neighbors=cfg.get("geo_token_k_neighbors",
-                                model.config.geo_token_k_neighbors),
-            dropout=cfg.get("geo_token_dropout",
-                            model.config.geo_token_dropout),
-            learning_rate=cfg.get(
-                "geo_token_learning_rate", model.config.geo_token_learning_rate),
-            epochs=int(cfg.get("geo_token_epochs",
-                       model.config.geo_token_epochs)),
-            patience=5,
-            task_type=model.task_type,
-            tweedie_power=tw_power,
-            use_data_parallel=False,
-            use_ddp=False,
-            use_approx_knn=model.config.gnn_use_approx_knn,
-            approx_knn_threshold=model.config.gnn_approx_knn_threshold,
-            graph_cache_path=None,
-            max_gpu_knn_nodes=model.config.gnn_max_gpu_knn_nodes,
-            knn_gpu_mem_ratio=model.config.gnn_knn_gpu_mem_ratio,
-            knn_gpu_mem_overhead=model.config.gnn_knn_gpu_mem_overhead
-        )
-        geo_gnn.fit(
-            train_geo,
-            model.train_data[model.resp_nme],
-            model.train_data[model.weight_nme]
-        )
-        train_embed = geo_gnn.encode(train_geo)
-        test_embed = geo_gnn.encode(test_geo)
-        cols = [f"geo_token_{i}" for i in range(train_embed.shape[1])]
-        train_tokens = pd.DataFrame(
-            train_embed, index=model.train_data.index, columns=cols)
-        test_tokens = pd.DataFrame(
-            test_embed, index=model.test_data.index, columns=cols)
-        return train_tokens, test_tokens, cols, geo_gnn
-    except Exception as exc:
-        print(f"[GeoToken] Generation failed: {exc}")
-        return None
-
-def _prepare_geo_tokens(model: "BayesOptModel") -> None:
-    """Build and persist geo tokens with default config values."""
-    gnn_trainer = model.model_manager.trainers.get("gnn")
-    if gnn_trainer is not None and hasattr(gnn_trainer, "prepare_geo_tokens"):
-        try:
-            gnn_trainer.prepare_geo_tokens(force=False)  # type: ignore[attr-defined]
-            return
-        except Exception as exc:
-            print(f"[GeoToken] GNNTrainer generation failed: {exc}")
-
-    result = _build_geo_tokens(model)
-    if result is None:
-        return
-    train_tokens, test_tokens, cols, geo_gnn = result
-    model.train_geo_tokens = train_tokens
-    model.test_geo_tokens = test_tokens
-    model.geo_token_cols = cols
-    model.geo_gnn_model = geo_gnn
-    print(f"[GeoToken] Generated {len(cols)}-dim geo tokens; injecting into FT.")

ins_pricing_gemini/modelling/model_manager.py

@@ -1,148 +0,0 @@
-from __future__ import annotations
-from typing import Dict, Optional, Any
-from datetime import datetime
-from dataclasses import asdict
-from .trainers import TrainerBase, GLMTrainer, XGBTrainer, ResNetTrainer, FTTrainer, GNNTrainer
-from .utils import IOUtils
-
-class ModelManager:
-    """Manages lifecycle and access to model trainers."""
-    
-    def __init__(self, context: Any) -> None:
-        # context is the BayesOptModel instance
-        self.ctx = context
-        self.trainers: Dict[str, TrainerBase] = {}
-        self._initialize_trainers()
-        
-    def _initialize_trainers(self) -> None:
-        self.trainers['glm'] = GLMTrainer(self.ctx)
-        self.trainers['xgb'] = XGBTrainer(self.ctx)
-        self.trainers['resn'] = ResNetTrainer(self.ctx)
-        self.trainers['ft'] = FTTrainer(self.ctx)
-        self.trainers['gnn'] = GNNTrainer(self.ctx)
-
-    def get_trainer(self, key: str) -> TrainerBase:
-        trainer = self.trainers.get(key)
-        if trainer is None:
-            raise KeyError(f"Unknown model key: {key}")
-        return trainer
-
-    def _maybe_load_best_params(self, model_key: str, trainer: TrainerBase) -> None:
-        # 1) If best_params_files is specified, load and skip tuning.
-        best_params_files = getattr(self.ctx.config, "best_params_files", None) or {}
-        best_params_file = best_params_files.get(model_key)
-        if best_params_file and not trainer.best_params:
-            trainer.best_params = IOUtils.load_params_file(best_params_file)
-            trainer.best_trial = None
-            print(
-                f"[Optuna][{trainer.label}] Loaded best_params from {best_params_file}; skip tuning."
-            )
-
-        # 2) If reuse_best_params is enabled, prefer version snapshots.
-        reuse_params = bool(getattr(self.ctx.config, "reuse_best_params", False))
-        if reuse_params and not trainer.best_params:
-            payload = self.ctx.version_manager.load_latest(f"{model_key}_best")
-            best_params = None if payload is None else payload.get("best_params")
-            if best_params:
-                trainer.best_params = best_params
-                trainer.best_trial = None
-                trainer.study_name = payload.get(
-                    "study_name") if isinstance(payload, dict) else None
-                print(
-                    f"[Optuna][{trainer.label}] Reusing best_params from versions snapshot.")
-                return
-            
-            # Fallback to legacy CSV (accessed via ctx.output_manager which is available on ctx)
-            params_path = self.ctx.output_manager.result_path(
-                f'{self.ctx.model_nme}_bestparams_{trainer.label.lower()}.csv'
-            )
-            # trainer.load_best_params_csv is not standard on TrainerBase but implemented on subclasses usually
-            # But checking core.py, it was loading locally.
-            # Actually core.py logic for (3) was omitted in my previous read view (lines 640+).
-            # Assuming I should rely on whatever logic was there or just omit legacy CSV if possible.
-            # But to be safe, let's stick to what we see: reusing snapshots is modern way.
-            # If logic requires CSV loading, I'd need to verify Trainer implementations. 
-            # Ideally Trainer.load() or similar should handle this?
-            # For now, I'll rely on version snapshots as primary persistence.
-
-    def optimize(self, model_key: str, max_evals: int = 100) -> None:
-        if model_key not in self.trainers:
-            print(f"Warning: Unknown model key: {model_key}")
-            return
-
-        trainer = self.get_trainer(model_key)
-        self._maybe_load_best_params(model_key, trainer)
-
-        should_tune = not trainer.best_params
-        if should_tune:
-            if model_key == "ft" and str(self.ctx.config.ft_role) == "unsupervised_embedding":
-                if hasattr(trainer, "cross_val_unsupervised"):
-                    trainer.tune(
-                        max_evals,
-                        objective_fn=getattr(trainer, "cross_val_unsupervised")
-                    )
-                else:
-                    raise RuntimeError(
-                        "FT trainer does not support unsupervised Optuna objective.")
-            else:
-                trainer.tune(max_evals)
-
-        if model_key == "ft" and str(self.ctx.config.ft_role) != "model":
-            prefix = str(self.ctx.config.ft_feature_prefix or "ft_emb")
-            role = str(self.ctx.config.ft_role)
-            if role == "embedding":
-                trainer.train_as_feature(
-                    pred_prefix=prefix, feature_mode="embedding")
-            elif role == "unsupervised_embedding":
-                trainer.pretrain_unsupervised_as_feature(
-                    pred_prefix=prefix,
-                    params=trainer.best_params
-                )
-            else:
-                raise ValueError(
-                    f"Unsupported ft_role='{role}', expected 'model'/'embedding'/'unsupervised_embedding'.")
-
-            # Inject generated prediction/embedding columns as features (scalar or vector).
-            # Callback to context since data state lives there (in DataContainer delegators)
-            self.ctx._inject_pred_features(prefix)
-        else:
-            trainer.train()
-
-        if bool(getattr(self.ctx.config, "final_ensemble", False)):
-            k = int(getattr(self.ctx.config, "final_ensemble_k", 3) or 3)
-            if k > 1:
-                if model_key == "ft" and str(self.ctx.config.ft_role) != "model":
-                    pass
-                elif hasattr(trainer, "ensemble_predict"):
-                    trainer.ensemble_predict(k)
-                else:
-                    print(
-                        f"[Ensemble] Trainer '{model_key}' does not support ensemble prediction.",
-                        flush=True,
-                    )
-
-        # Update context fields for backward compatibility
-        setattr(self.ctx, f"{model_key}_best", trainer.model)
-        setattr(self.ctx, f"best_{model_key}_params", trainer.best_params)
-        setattr(self.ctx, f"best_{model_key}_trial", trainer.best_trial)
-        
-        # Save a snapshot for traceability
-        study_name = getattr(trainer, "study_name", None)
-        if study_name is None and trainer.best_trial is not None:
-            study_obj = getattr(trainer.best_trial, "study", None)
-            study_name = getattr(study_obj, "study_name", None)
-            
-        # Pydantic config to dict
-        if hasattr(self.ctx.config, "model_dump"):
-            config_dict = self.ctx.config.model_dump()
-        else:
-            config_dict = asdict(self.ctx.config) # Fallback if for some reason it's dataclass (shouldn't be)
-
-        snapshot = {
-            "model_key": model_key,
-            "timestamp": datetime.now().isoformat(),
-            "best_params": trainer.best_params,
-            "study_name": study_name,
-            "config": config_dict,
-        }
-        self.ctx.version_manager.save(f"{model_key}_best", snapshot)