Add causal_selection/meta_learner/trainer.py

6f65cd1 verified 7 months ago

12.4 kB

"""
Meta-learner: trains models to predict algorithm performance from dataset meta-features.
Supports multi-output regression (predict SHD per algorithm) and ranking evaluation.
"""
import os
import numpy as np
import pandas as pd
import json
import logging
from collections import defaultdict

from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.multioutput import MultiOutputRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import LeaveOneGroupOut, cross_val_predict
from sklearn.metrics import mean_squared_error, mean_absolute_error
import joblib

from causal_selection.features.extractor import FEATURE_NAMES
from causal_selection.discovery.algorithms import ALGORITHM_POOL

logger = logging.getLogger(__name__)

ALGO_NAMES = list(ALGORITHM_POOL.keys())
RESULTS_DIR = '/app/causal_selection/data/results'
MODEL_DIR = '/app/causal_selection/models'


def load_meta_dataset(results_dir=RESULTS_DIR):
    """Load meta-dataset from CSV files."""
    X = pd.read_csv(os.path.join(results_dir, 'meta_features.csv'))
    Y_shd = pd.read_csv(os.path.join(results_dir, 'shd_matrix.csv'))
    Y_nshd = pd.read_csv(os.path.join(results_dir, 'normalized_shd_matrix.csv'))
    configs = pd.read_csv(os.path.join(results_dir, 'configs.csv'))
    return X, Y_shd, Y_nshd, configs


def train_meta_learner(X, Y, model_type='rf', **model_kwargs):
    """Train a multi-output regression model.
    
    Args:
        X: feature matrix (n_tasks, n_features)
        Y: target matrix (n_tasks, n_algorithms) - SHD values
        model_type: 'rf' or 'gbm'
    
    Returns:
        trained model, scaler
    """
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)
    
    if model_type == 'rf':
        base = RandomForestRegressor(
            n_estimators=model_kwargs.get('n_estimators', 200),
            max_depth=model_kwargs.get('max_depth', None),
            min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
            random_state=42,
            n_jobs=-1,
        )
    elif model_type == 'gbm':
        base = GradientBoostingRegressor(
            n_estimators=model_kwargs.get('n_estimators', 200),
            max_depth=model_kwargs.get('max_depth', 5),
            learning_rate=model_kwargs.get('learning_rate', 0.1),
            min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
            random_state=42,
        )
    else:
        raise ValueError(f"Unknown model type: {model_type}")
    
    model = MultiOutputRegressor(base)
    model.fit(X_scaled, Y)
    
    return model, scaler


def predict_top_k(model, scaler, X_new, k=3):
    """Predict top-k algorithms for new dataset(s).
    
    Args:
        model: trained multi-output model
        scaler: fitted StandardScaler
        X_new: feature matrix (n_new, n_features)
        k: number of top algorithms to return
    
    Returns:
        top_k_indices: (n_new, k) array of algorithm indices (sorted by predicted SHD ascending)
        predicted_shd: (n_new, n_algorithms) full predicted SHD matrix
    """
    X_scaled = scaler.transform(X_new)
    predicted = model.predict(X_scaled)
    
    if predicted.ndim == 1:
        predicted = predicted.reshape(1, -1)
    
    top_k_indices = np.argsort(predicted, axis=1)[:, :k]
    return top_k_indices, predicted


def evaluate_lono_cv(X, Y, configs, model_type='rf', k=3, **model_kwargs):
    """Leave-One-Network-Out Cross-Validation.
    
    For each network, train on all other networks, test on that network.
    This tests generalization to truly unseen graph structures.
    
    Returns:
        results: dict with metrics per network and overall
    """
    networks = configs['network'].values
    unique_networks = sorted(configs['network'].unique())
    
    results = {
        'per_network': {},
        'all_predictions': [],
        'all_true': [],
        'all_configs': [],
    }
    
    scaler = StandardScaler()
    
    for test_net in unique_networks:
        test_mask = networks == test_net
        train_mask = ~test_mask
        
        if train_mask.sum() < 3:
            logger.warning(f"Skipping {test_net}: only {train_mask.sum()} training samples")
            continue
        
        X_train = X.values[train_mask]
        Y_train = Y.values[train_mask]
        X_test = X.values[test_mask]
        Y_test = Y.values[test_mask]
        
        # Scale
        scaler.fit(X_train)
        X_train_s = scaler.transform(X_train)
        X_test_s = scaler.transform(X_test)
        
        # Train
        if model_type == 'rf':
            base = RandomForestRegressor(
                n_estimators=model_kwargs.get('n_estimators', 200),
                max_depth=model_kwargs.get('max_depth', None),
                min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
                random_state=42, n_jobs=-1,
            )
        else:
            base = GradientBoostingRegressor(
                n_estimators=model_kwargs.get('n_estimators', 200),
                max_depth=model_kwargs.get('max_depth', 5),
                learning_rate=model_kwargs.get('learning_rate', 0.1),
                min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
                random_state=42,
            )
        
        model = MultiOutputRegressor(base)
        model.fit(X_train_s, Y_train)
        
        # Predict
        Y_pred = model.predict(X_test_s)
        
        # Evaluate
        net_metrics = _compute_ranking_metrics(Y_pred, Y_test, k=k)
        net_metrics['n_test'] = int(test_mask.sum())
        net_metrics['n_train'] = int(train_mask.sum())
        results['per_network'][test_net] = net_metrics
        
        results['all_predictions'].extend(Y_pred.tolist())
        results['all_true'].extend(Y_test.tolist())
        results['all_configs'].extend(
            configs[test_mask][['network', 'n_samples', 'seed']].to_dict('records')
        )
        
        logger.info(f"  {test_net:15s}: top{k}_hit={net_metrics['top_k_hit_rate']:.3f} "
                    f"regret={net_metrics['mean_regret']:.2f} "
                    f"ndcg={net_metrics['ndcg_at_k']:.3f}")
    
    # Overall metrics
    all_pred = np.array(results['all_predictions'])
    all_true = np.array(results['all_true'])
    overall = _compute_ranking_metrics(all_pred, all_true, k=k)
    results['overall'] = overall
    
    return results


def _compute_ranking_metrics(Y_pred, Y_true, k=3):
    """Compute ranking metrics for algorithm selection.
    
    Args:
        Y_pred: (n, n_algos) predicted SHD values
        Y_true: (n, n_algos) true SHD values
        k: top-k to consider
    """
    n = Y_pred.shape[0]
    
    top_k_hits = 0
    regrets = []
    ndcgs = []
    
    for i in range(n):
        true_ranking = np.argsort(Y_true[i])  # best algo first
        pred_ranking = np.argsort(Y_pred[i])  # predicted best first
        
        true_best = true_ranking[0]
        pred_top_k = pred_ranking[:k]
        
        # Top-k hit rate: is the true best in predicted top-k?
        if true_best in pred_top_k:
            top_k_hits += 1
        
        # SHD regret: SHD of best in predicted top-k minus oracle best SHD
        oracle_shd = Y_true[i, true_best]
        selected_shds = [Y_true[i, j] for j in pred_top_k]
        best_selected_shd = min(selected_shds)
        regret = best_selected_shd - oracle_shd
        regrets.append(regret)
        
        # NDCG@k
        ndcg = _ndcg_at_k(Y_true[i], Y_pred[i], k)
        ndcgs.append(ndcg)
    
    # Also compute: is one of the true top-3 in the predicted top-3?
    top_k_overlap = 0
    for i in range(n):
        true_top_k = set(np.argsort(Y_true[i])[:k])
        pred_top_k = set(np.argsort(Y_pred[i])[:k])
        overlap = len(true_top_k & pred_top_k)
        top_k_overlap += overlap / k
    
    return {
        'top_k_hit_rate': top_k_hits / n,  # true best in predicted top-k
        'top_k_overlap_rate': top_k_overlap / n,  # avg overlap between true/pred top-k
        'mean_regret': np.mean(regrets),
        'median_regret': np.median(regrets),
        'max_regret': np.max(regrets),
        'ndcg_at_k': np.mean(ndcgs),
        'mean_pred_mse': mean_squared_error(Y_true, Y_pred),
        'mean_pred_mae': mean_absolute_error(Y_true, Y_pred),
    }


def _ndcg_at_k(true_scores, pred_scores, k):
    """Normalized Discounted Cumulative Gain at k.
    
    For algorithm selection: lower SHD = better, so we negate scores for ranking.
    """
    # Convert SHD to relevance: rel = max_shd - shd (higher = better)
    max_shd = max(true_scores.max(), 1)
    relevance = max_shd - true_scores
    
    # Predicted ranking
    pred_order = np.argsort(pred_scores)[:k]
    
    # DCG
    dcg = 0
    for rank, idx in enumerate(pred_order):
        dcg += relevance[idx] / np.log2(rank + 2)
    
    # Ideal DCG
    ideal_order = np.argsort(-relevance)[:k]
    idcg = 0
    for rank, idx in enumerate(ideal_order):
        idcg += relevance[idx] / np.log2(rank + 2)
    
    return dcg / idcg if idcg > 0 else 0


def get_feature_importance(model, feature_names=FEATURE_NAMES, algo_names=ALGO_NAMES):
    """Extract feature importance from trained model."""
    importances = {}
    for i, (algo, estimator) in enumerate(zip(algo_names, model.estimators_)):
        if hasattr(estimator, 'feature_importances_'):
            importances[algo] = dict(zip(feature_names, estimator.feature_importances_))
    
    # Average importance across algorithms
    avg_importance = defaultdict(float)
    for algo, imp in importances.items():
        for feat, val in imp.items():
            avg_importance[feat] += val / len(importances)
    
    return dict(avg_importance), importances


def save_model(model, scaler, model_dir=MODEL_DIR):
    """Save trained model and scaler."""
    os.makedirs(model_dir, exist_ok=True)
    joblib.dump(model, os.path.join(model_dir, 'meta_learner.pkl'))
    joblib.dump(scaler, os.path.join(model_dir, 'scaler.pkl'))
    logger.info(f"Model saved to {model_dir}")


def load_model(model_dir=MODEL_DIR):
    """Load trained model and scaler."""
    model = joblib.load(os.path.join(model_dir, 'meta_learner.pkl'))
    scaler = joblib.load(os.path.join(model_dir, 'scaler.pkl'))
    return model, scaler


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
    
    # Load meta-dataset
    X, Y_shd, Y_nshd, configs = load_meta_dataset()
    
    print(f"Meta-dataset: X={X.shape}, Y_shd={Y_shd.shape}")
    print(f"Networks: {configs['network'].unique()}")
    print(f"Configs per network:")
    print(configs['network'].value_counts().to_string())
    
    # Evaluate with LONO-CV
    print("\n" + "=" * 80)
    print("LEAVE-ONE-NETWORK-OUT CV (RandomForest)")
    print("=" * 80)
    
    results_rf = evaluate_lono_cv(X, Y_nshd, configs, model_type='rf', k=3)
    
    print(f"\nOverall Results (RF):")
    for k, v in results_rf['overall'].items():
        print(f"  {k:25s}: {v:.4f}")
    
    print("\n" + "=" * 80)
    print("LEAVE-ONE-NETWORK-OUT CV (GradientBoosting)")
    print("=" * 80)
    
    results_gbm = evaluate_lono_cv(X, Y_nshd, configs, model_type='gbm', k=3)
    
    print(f"\nOverall Results (GBM):")
    for k, v in results_gbm['overall'].items():
        print(f"  {k:25s}: {v:.4f}")
    
    # Train final model on all data
    print("\n" + "=" * 80)
    print("TRAINING FINAL MODEL")
    print("=" * 80)
    
    best_type = 'rf' if results_rf['overall']['top_k_hit_rate'] >= results_gbm['overall']['top_k_hit_rate'] else 'gbm'
    print(f"Selected model type: {best_type}")
    
    model, scaler = train_meta_learner(X, Y_nshd, model_type=best_type)
    save_model(model, scaler)
    
    # Feature importance
    avg_imp, per_algo_imp = get_feature_importance(model)
    print("\nTop 10 Most Important Features:")
    for feat, imp in sorted(avg_imp.items(), key=lambda x: -x[1])[:10]:
        print(f"  {feat:30s}: {imp:.4f}")
    
    # Save all evaluation results
    with open(os.path.join(RESULTS_DIR, 'evaluation_results.json'), 'w') as f:
        json.dump({
            'rf': {k: v for k, v in results_rf['overall'].items()},
            'gbm': {k: v for k, v in results_gbm['overall'].items()},
            'feature_importance': avg_imp,
            'selected_model': best_type,
        }, f, indent=2)

	"""
	Meta-learner: trains models to predict algorithm performance from dataset meta-features.
	Supports multi-output regression (predict SHD per algorithm) and ranking evaluation.
	"""
	import os
	import numpy as np
	import pandas as pd
	import json
	import logging
	from collections import defaultdict

	from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
	from sklearn.multioutput import MultiOutputRegressor
	from sklearn.preprocessing import StandardScaler
	from sklearn.model_selection import LeaveOneGroupOut, cross_val_predict
	from sklearn.metrics import mean_squared_error, mean_absolute_error
	import joblib

	from causal_selection.features.extractor import FEATURE_NAMES
	from causal_selection.discovery.algorithms import ALGORITHM_POOL

	logger = logging.getLogger(__name__)

	ALGO_NAMES = list(ALGORITHM_POOL.keys())
	RESULTS_DIR = '/app/causal_selection/data/results'
	MODEL_DIR = '/app/causal_selection/models'


	def load_meta_dataset(results_dir=RESULTS_DIR):
	"""Load meta-dataset from CSV files."""
	X = pd.read_csv(os.path.join(results_dir, 'meta_features.csv'))
	Y_shd = pd.read_csv(os.path.join(results_dir, 'shd_matrix.csv'))
	Y_nshd = pd.read_csv(os.path.join(results_dir, 'normalized_shd_matrix.csv'))
	configs = pd.read_csv(os.path.join(results_dir, 'configs.csv'))
	return X, Y_shd, Y_nshd, configs


	def train_meta_learner(X, Y, model_type='rf', **model_kwargs):
	"""Train a multi-output regression model.

	Args:
	X: feature matrix (n_tasks, n_features)
	Y: target matrix (n_tasks, n_algorithms) - SHD values
	model_type: 'rf' or 'gbm'

	Returns:
	trained model, scaler
	"""
	scaler = StandardScaler()
	X_scaled = scaler.fit_transform(X)

	if model_type == 'rf':
	base = RandomForestRegressor(
	n_estimators=model_kwargs.get('n_estimators', 200),
	max_depth=model_kwargs.get('max_depth', None),
	min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
	random_state=42,
	n_jobs=-1,
	)
	elif model_type == 'gbm':
	base = GradientBoostingRegressor(
	n_estimators=model_kwargs.get('n_estimators', 200),
	max_depth=model_kwargs.get('max_depth', 5),
	learning_rate=model_kwargs.get('learning_rate', 0.1),
	min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
	random_state=42,
	)
	else:
	raise ValueError(f"Unknown model type: {model_type}")

	model = MultiOutputRegressor(base)
	model.fit(X_scaled, Y)

	return model, scaler


	def predict_top_k(model, scaler, X_new, k=3):
	"""Predict top-k algorithms for new dataset(s).

	Args:
	model: trained multi-output model
	scaler: fitted StandardScaler
	X_new: feature matrix (n_new, n_features)
	k: number of top algorithms to return

	Returns:
	top_k_indices: (n_new, k) array of algorithm indices (sorted by predicted SHD ascending)
	predicted_shd: (n_new, n_algorithms) full predicted SHD matrix
	"""
	X_scaled = scaler.transform(X_new)
	predicted = model.predict(X_scaled)

	if predicted.ndim == 1:
	predicted = predicted.reshape(1, -1)

	top_k_indices = np.argsort(predicted, axis=1)[:, :k]
	return top_k_indices, predicted


	def evaluate_lono_cv(X, Y, configs, model_type='rf', k=3, **model_kwargs):
	"""Leave-One-Network-Out Cross-Validation.

	For each network, train on all other networks, test on that network.
	This tests generalization to truly unseen graph structures.

	Returns:
	results: dict with metrics per network and overall
	"""
	networks = configs['network'].values
	unique_networks = sorted(configs['network'].unique())

	results = {
	'per_network': {},
	'all_predictions': [],
	'all_true': [],
	'all_configs': [],
	}

	scaler = StandardScaler()

	for test_net in unique_networks:
	test_mask = networks == test_net
	train_mask = ~test_mask

	if train_mask.sum() < 3:
	logger.warning(f"Skipping {test_net}: only {train_mask.sum()} training samples")
	continue

	X_train = X.values[train_mask]
	Y_train = Y.values[train_mask]
	X_test = X.values[test_mask]
	Y_test = Y.values[test_mask]

	# Scale
	scaler.fit(X_train)
	X_train_s = scaler.transform(X_train)
	X_test_s = scaler.transform(X_test)

	# Train
	if model_type == 'rf':
	base = RandomForestRegressor(
	n_estimators=model_kwargs.get('n_estimators', 200),
	max_depth=model_kwargs.get('max_depth', None),
	min_samples_leaf=model_kwargs.get('min_samples_leaf', 2),
	random_state=42, n_jobs=-1,
	)
	else:
	base = GradientBoostingRegressor(
	n_estimators=model_kwargs.get('n_estimators', 200),
	max_depth=model_kwargs.get('max_depth', 5),
	learning_rate=model_kwargs.get('learning_rate', 0.1),
	min_samples_leaf=model_kwargs.get('min_samples_leaf', 3),
	random_state=42,
	)

	model = MultiOutputRegressor(base)
	model.fit(X_train_s, Y_train)

	# Predict
	Y_pred = model.predict(X_test_s)

	# Evaluate
	net_metrics = _compute_ranking_metrics(Y_pred, Y_test, k=k)
	net_metrics['n_test'] = int(test_mask.sum())
	net_metrics['n_train'] = int(train_mask.sum())
	results['per_network'][test_net] = net_metrics

	results['all_predictions'].extend(Y_pred.tolist())
	results['all_true'].extend(Y_test.tolist())
	results['all_configs'].extend(
	configs[test_mask][['network', 'n_samples', 'seed']].to_dict('records')
	)

	logger.info(f" {test_net:15s}: top{k}_hit={net_metrics['top_k_hit_rate']:.3f} "
	f"regret={net_metrics['mean_regret']:.2f} "
	f"ndcg={net_metrics['ndcg_at_k']:.3f}")

	# Overall metrics
	all_pred = np.array(results['all_predictions'])
	all_true = np.array(results['all_true'])
	overall = _compute_ranking_metrics(all_pred, all_true, k=k)
	results['overall'] = overall

	return results


	def _compute_ranking_metrics(Y_pred, Y_true, k=3):
	"""Compute ranking metrics for algorithm selection.

	Args:
	Y_pred: (n, n_algos) predicted SHD values
	Y_true: (n, n_algos) true SHD values
	k: top-k to consider
	"""
	n = Y_pred.shape[0]

	top_k_hits = 0
	regrets = []
	ndcgs = []

	for i in range(n):
	true_ranking = np.argsort(Y_true[i]) # best algo first
	pred_ranking = np.argsort(Y_pred[i]) # predicted best first

	true_best = true_ranking[0]
	pred_top_k = pred_ranking[:k]

	# Top-k hit rate: is the true best in predicted top-k?
	if true_best in pred_top_k:
	top_k_hits += 1

	# SHD regret: SHD of best in predicted top-k minus oracle best SHD
	oracle_shd = Y_true[i, true_best]
	selected_shds = [Y_true[i, j] for j in pred_top_k]
	best_selected_shd = min(selected_shds)
	regret = best_selected_shd - oracle_shd
	regrets.append(regret)

	# NDCG@k
	ndcg = _ndcg_at_k(Y_true[i], Y_pred[i], k)
	ndcgs.append(ndcg)

	# Also compute: is one of the true top-3 in the predicted top-3?
	top_k_overlap = 0
	for i in range(n):
	true_top_k = set(np.argsort(Y_true[i])[:k])
	pred_top_k = set(np.argsort(Y_pred[i])[:k])
	overlap = len(true_top_k & pred_top_k)
	top_k_overlap += overlap / k

	return {
	'top_k_hit_rate': top_k_hits / n, # true best in predicted top-k
	'top_k_overlap_rate': top_k_overlap / n, # avg overlap between true/pred top-k
	'mean_regret': np.mean(regrets),
	'median_regret': np.median(regrets),
	'max_regret': np.max(regrets),
	'ndcg_at_k': np.mean(ndcgs),
	'mean_pred_mse': mean_squared_error(Y_true, Y_pred),
	'mean_pred_mae': mean_absolute_error(Y_true, Y_pred),
	}


	def _ndcg_at_k(true_scores, pred_scores, k):
	"""Normalized Discounted Cumulative Gain at k.

	For algorithm selection: lower SHD = better, so we negate scores for ranking.
	"""
	# Convert SHD to relevance: rel = max_shd - shd (higher = better)
	max_shd = max(true_scores.max(), 1)
	relevance = max_shd - true_scores

	# Predicted ranking
	pred_order = np.argsort(pred_scores)[:k]

	# DCG
	dcg = 0
	for rank, idx in enumerate(pred_order):
	dcg += relevance[idx] / np.log2(rank + 2)

	# Ideal DCG
	ideal_order = np.argsort(-relevance)[:k]
	idcg = 0
	for rank, idx in enumerate(ideal_order):
	idcg += relevance[idx] / np.log2(rank + 2)

	return dcg / idcg if idcg > 0 else 0


	def get_feature_importance(model, feature_names=FEATURE_NAMES, algo_names=ALGO_NAMES):
	"""Extract feature importance from trained model."""
	importances = {}
	for i, (algo, estimator) in enumerate(zip(algo_names, model.estimators_)):
	if hasattr(estimator, 'feature_importances_'):
	importances[algo] = dict(zip(feature_names, estimator.feature_importances_))

	# Average importance across algorithms
	avg_importance = defaultdict(float)
	for algo, imp in importances.items():
	for feat, val in imp.items():
	avg_importance[feat] += val / len(importances)

	return dict(avg_importance), importances


	def save_model(model, scaler, model_dir=MODEL_DIR):
	"""Save trained model and scaler."""
	os.makedirs(model_dir, exist_ok=True)
	joblib.dump(model, os.path.join(model_dir, 'meta_learner.pkl'))
	joblib.dump(scaler, os.path.join(model_dir, 'scaler.pkl'))
	logger.info(f"Model saved to {model_dir}")


	def load_model(model_dir=MODEL_DIR):
	"""Load trained model and scaler."""
	model = joblib.load(os.path.join(model_dir, 'meta_learner.pkl'))
	scaler = joblib.load(os.path.join(model_dir, 'scaler.pkl'))
	return model, scaler


	if __name__ == '__main__':
	logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')

	# Load meta-dataset
	X, Y_shd, Y_nshd, configs = load_meta_dataset()

	print(f"Meta-dataset: X={X.shape}, Y_shd={Y_shd.shape}")
	print(f"Networks: {configs['network'].unique()}")
	print(f"Configs per network:")
	print(configs['network'].value_counts().to_string())

	# Evaluate with LONO-CV
	print("\n" + "=" * 80)
	print("LEAVE-ONE-NETWORK-OUT CV (RandomForest)")
	print("=" * 80)

	results_rf = evaluate_lono_cv(X, Y_nshd, configs, model_type='rf', k=3)

	print(f"\nOverall Results (RF):")
	for k, v in results_rf['overall'].items():
	print(f" {k:25s}: {v:.4f}")

	print("\n" + "=" * 80)
	print("LEAVE-ONE-NETWORK-OUT CV (GradientBoosting)")
	print("=" * 80)

	results_gbm = evaluate_lono_cv(X, Y_nshd, configs, model_type='gbm', k=3)

	print(f"\nOverall Results (GBM):")
	for k, v in results_gbm['overall'].items():
	print(f" {k:25s}: {v:.4f}")

	# Train final model on all data
	print("\n" + "=" * 80)
	print("TRAINING FINAL MODEL")
	print("=" * 80)

	best_type = 'rf' if results_rf['overall']['top_k_hit_rate'] >= results_gbm['overall']['top_k_hit_rate'] else 'gbm'
	print(f"Selected model type: {best_type}")

	model, scaler = train_meta_learner(X, Y_nshd, model_type=best_type)
	save_model(model, scaler)

	# Feature importance
	avg_imp, per_algo_imp = get_feature_importance(model)
	print("\nTop 10 Most Important Features:")
	for feat, imp in sorted(avg_imp.items(), key=lambda x: -x[1])[:10]:
	print(f" {feat:30s}: {imp:.4f}")

	# Save all evaluation results
	with open(os.path.join(RESULTS_DIR, 'evaluation_results.json'), 'w') as f:
	json.dump({
	'rf': {k: v for k, v in results_rf['overall'].items()},
	'gbm': {k: v for k, v in results_gbm['overall'].items()},
	'feature_importance': avg_imp,
	'selected_model': best_type,
	}, f, indent=2)