<reponame>wikimedia/search-MjoLniR<gh_stars>10-100
from collections import defaultdict
from functools import partial
import tempfile
from typing import cast, Any, Callable, Dict, List, Mapping, Optional

import hyperopt
import numpy as np
from pyspark.sql import SparkSession
import xgboost as xgb

from mjolnir.training.tuning import make_cv_objective, ModelSelection
from mjolnir.utils import as_local_path, as_local_paths, as_output_file


def _coerce_params(params: Mapping[str, Any]) -> Dict[str, Any]:
    """Force xgboost parameters into appropriate types

    The output from hyperopt is always floats, but some xgboost parameters
    explicitly require integers. Cast those as necessary

    Parameters
    ----------
    params : dict
        xgboost parameters

    Returns
    -------
    dict
        Input parameters coerced as necessary
    """
    def identity(x):
        return x

    def sloppy_int(x):
        try:
            return int(x)
        except ValueError:
            pass

        val = float(x)
        # This could fail for larger numbers due to fp precision, but not
        # expecting integer values larger than two digits here.
        if val.is_integer():
            return int(val)
        raise ValueError('Not parsable as integer: {}'.format(x))

    types = cast(Dict[str, Callable[[Any], Any]], defaultdict(lambda: identity))
    types.update({
        'max_depth': sloppy_int,
        'max_bin': sloppy_int,
        'num_class': sloppy_int,
        'silent': sloppy_int,
    })
    return {k: types[k](v) for k, v in params.items()}


def train(
    fold: Mapping[str, str],
    params: Mapping[str, Any],
    train_matrix: Optional[str] = None,
    spark: Optional[SparkSession] = None
) -> 'XGBoostModel':
    """Train a single xgboost ranking model.

    Primary entry point for hyperparameter tuning normalizes
    parameters and auto detects some values. Actual training is
    passed on to XGBoostModel.trainWithFiles

    Parameters
    ----------
    fold :
        Map from split name to data path. All provided splits will be
        evaluated on each boosting iteration.
    params :
        parameters to pass on to xgboost training
    train_matrix :
        Optional name of training matrix in fold. If not provided will
        auto-detect to either 'all' or 'train'
    spark:
        If provided, train remotely over spark

    Returns
    -------
    XGBoostModel
        Trained xgboost model
    """
    # hyperparameter tuning may have given us floats where we need
    # ints, so this gets all the types right for Java. Also makes
    # a copy of params so we don't modifying the incoming dict.
    # TODO: Does python care about this?  Even if not, in the end it seems
    # reasonable to not pass floats for integer values.
    params = _coerce_params(params)

    # TODO: Maybe num_rounds should just be external? But it's easier
    # to do hyperparameter optimization with a consistent dict interface
    kwargs = cast(Dict[str, Any], {
        'num_boost_round': 100,
    })

    if 'num_boost_round' in params:
        kwargs['num_boost_round'] = params['num_boost_round']
        del params['num_rounds']
    if 'early_stopping_rounds' in params:
        kwargs['early_stopping_rounds'] = params['early_stopping_rounds']
        del params['early_stopping_rounds']

    # Set some sane defaults for ranking tasks
    if 'objective' not in params:
        params['objective'] = 'rank:ndcg'
    if 'eval_metric' not in params:
        params['eval_metric'] = 'ndcg@10'
    # Not really ranking specific, but generally fastest method
    if 'tree_method' not in params:
        params['tree_method'] = 'hist'
    # Convenience for some situations, but typically be explicit
    # about the name of the matrix to train against.
    if train_matrix is None:
        train_matrix = "all" if "all" in fold else "train"

    if spark:
        return XGBoostModel.trainWithFilesRemote(spark, fold, train_matrix, params, **kwargs)
    else:
        return XGBoostModel.trainWithFiles(fold, train_matrix, params, **kwargs)


# Top level: matrix name
# Second level: metric name
# Inner list: stringified per-iteration metric value
EvalsResult = Mapping[str, Mapping[str, List[str]]]


class XGBoostBooster(object):
    """Wrapper for xgb.Booster usage in mjolnir

    Wraps the booster to distinguish what we have after training,
    the XGBoostModel, from what we write to disk, which is only the
    booster. Would be better if there was a clean way to wrap all
    the data up an serialize together, while working with xgboost's
    c++ methods that expect file paths.
    """
    def __init__(self, booster: xgb.Booster) -> None:
        self.booster = booster

    @staticmethod
    def loadBoosterFromHadoopFile(path: str) -> 'XGBoostBooster':
        with as_local_path(path) as local_path:
            return XGBoostBooster.loadBoosterFromLocalFile(local_path)

    @staticmethod
    def loadBoosterFromLocalFile(path: str) -> 'XGBoostBooster':
        booster = xgb.Booster.load_model(path)
        # TODO: Not having the training parameters or the evaluation metrics
        # almost makes this a different thing...
        return XGBoostBooster(booster)

    def saveBoosterAsHadoopFile(self, path: str):
        with as_output_file(path) as f:
            self.saveBoosterAsLocalFile(f.name)

    def saveBoosterAsLocalFile(self, path: str):
        # TODO: This doesn't save any metrics, should it?
        self.booster.save_model(path)


class XGBoostModel(XGBoostBooster):
    """xgboost booster along with train-time metrics

    TODO: Take XGBoostBooster as init arg instead of xgb.Booster?
    """

    def __init__(
        self,
        booster: xgb.Booster,
        evals_result: EvalsResult,
        params: Mapping[str, Any]
    ) -> None:
        super().__init__(booster)
        self.evals_result = evals_result
        self.params = params

    @staticmethod
    def trainWithFilesRemote(
        spark: SparkSession,
        fold: Mapping[str, str],
        train_matrix: str,
        params: Mapping[str, Any],
        **kwargs
    ) -> 'XGBoostModel':
        """Train model on a single remote spark executor.

        Silly hack to train models inside the yarn cluster. To train multiple
        models in parallel python threads will need to be used.  Wish pyspark
        had collectAsync.
        """
        nthread = int(spark.conf.get('spark.task.cpus', '1'))
        if 'nthread' not in params:
            params = dict(params, nthread=nthread)
        elif params['nthread'] != nthread:
            raise Exception("Executors have [{}] cpus but training requested [{}]".format(
                nthread, params['nthread']))

        return (
            spark.sparkContext
            .parallelize([1], 1)
            .map(lambda x: XGBoostModel.trainWithFiles(fold, train_matrix, params, **kwargs))
            .collect()[0]
        )

    @staticmethod
    def trainWithFiles(
        fold: Mapping[str, str],
        train_matrix: str,
        params: Mapping[str, Any],
        **kwargs
    ) -> 'XGBoostModel':
        """Wrapper around xgb.train

        This intentionally forwards to trainWithRDD, rather than
        trainWithDataFrame, as the underlying method currently prevents using
        rank:pairwise and metrics with @, such as ndcg@5.

        Parameters
        ----------
        fold :
            Map from split name to data path. All provided splits will be
            evaluated on each boosting iteration.
        train_matrix: str
            name of split in fold to train against
        params : dict
            XGBoost training parameters

        Returns
        -------
        mjolnir.training.xgboost.XGBoostModel
            trained xgboost ranking model
        """
        with as_local_paths(fold.values()) as local_paths:
            matrices = {name: xgb.DMatrix(path) for name, path in zip(fold.keys(), local_paths)}
            dtrain = matrices[train_matrix]
            evallist = [(dmat, name) for name, dmat in matrices.items()]
            metrics = cast(Mapping, {})
            booster = xgb.train(params, dtrain, evals=evallist, evals_result=metrics, **kwargs)
            return XGBoostModel(booster, metrics, params)

    def dump(self, features=None, with_stats=False, format="json"):
        """Dumps the xgboost model

        Parameters
        ----------
        features : list of str or None, optional
            list of features names, or None for no feature names in dump.
            (Default: None)
        withStats : bool, optional
            Should various additional statistics be included? These are not
            necessary for prediction. (Default: False)
        format : string, optional
            The format of dump to produce, either json or text. (Default: json)

        Returns
        -------
        str
            valid json string containing all trees
        """
        # Annoyingly the xgboost api doesn't take the feature map as a string, but
        # instead as a filename. Write the feature map out to a file if necessary.
        if features:
            # When we write the svmrank formatted files the features are indexed
            # starting at 1. We need to throw a fake index 0 in here or it's all
            # wrong.
            feat_map = "0 PLACEHOLDER_FEAT q\n" + \
                "\n".join(["%d %s q" % (i, fname) for i, fname in enumerate(features, 1)])
            fmap_f = tempfile.NamedTemporaryFile(mode='w')
            fmap_f.write(feat_map)
            fmap_f.flush()
            fmap_path = fmap_f.name
        else:
            fmap_path = ''

        trees = self.booster.get_dump(fmap_path, with_stats, dump_format='json')
        # For whatever reason we get a json line per tree. Turn that into an array
        # so we have a single valid json string.
        return '[' + ','.join(trees) + ']'


def cv_transformer(model: XGBoostModel, params: Mapping[str, Any]):
    """Report model metrics in format expected by model selection"""
    metric = params['eval_metric']
    return {
        'train': model.evals_result['train'][metric][-1],
        'test': model.evals_result['test'][metric][-1],
        'metrics': model.evals_result,
    }


def tune(
    folds: List[Mapping[str, str]],
    stats: Dict,
    train_matrix: str,
    num_cv_jobs: int = 5,
    initial_num_trees: int = 100,
    final_num_trees: int = 500,
    iterations: int = 150,
    spark: Optional[SparkSession] = None
):
    """Find appropriate hyperparameters for training df

    This is far from perfect, hyperparameter tuning is a bit of a black art
    and could probably benefit from human interaction at each stage. Various
    parameters depend a good bit on the number of samples in df, and how
    that data is shaped.

    Below is tuned for a dataframe with approximatly 10k normalized queries,
    110k total queries, and 2.2M samples. This is actually a relatively small
    dataset, we should rework the values used with larger data sets if they
    are promising. It may also be that the current feature space can't take
    advantage of more samples.

    Note that hyperopt uses the first 20 iterations to initialize, during those
    first 20 this is a strictly random search.

    Parameters
    ----------
    folds : list of dict containing train and test keys
    stats : dict
        stats about the fold from the make_folds utility script
    num_cv_jobs : int, optional
        The number of cross validation folds to train in parallel. (Default: 5)
    initial_num_trees: int, optional
        The number of trees to do most of the hyperparameter tuning with. This
        should be large enough to be resonably representative of the final
        training size. (Default: 100)
    final_num_trees: int, optional
        The number of trees to do the final eta optimization with. If set to
        None the final eta optimization will be skipped and initial_n_tree will
        be kept.

    Returns
    -------
    dict
        Dict with two keys, trials and params. params is the optimal set of
        parameters. trials contains a dict of individual optimization steps
        performed, each containing a hyperopt.Trials object recording what
        happened.
    """
    num_obs = stats['num_observations']

    if num_obs > 8000000:
        dataset_size = 'xlarge'
    elif num_obs > 1000000:
        dataset_size = 'large'
    elif num_obs > 500000:
        dataset_size = 'med'
    elif num_obs > 500:
        dataset_size = 'small'
    else:
        dataset_size = 'xsmall'

    # Setup different tuning profiles for different sizes of datasets.
    tune_spaces = [
        ('initial', {
            'iterations': iterations,
            'space': {
                'xlarge': {
                    'eta': hyperopt.hp.uniform('eta', 0.3, 0.8),
                    # Have seen values of 7 and 10 as best on roughly same size
                    # datasets from different wikis. It really just depends.
                    'max_depth': hyperopt.hp.quniform('max_depth', 6, 11, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(500), 10),
                    # % of features to use for each tree. helps prevent overfit
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'large': {
                    'eta': hyperopt.hp.uniform('eta', 0.3, 0.6),
                    'max_depth': hyperopt.hp.quniform('max_depth', 5, 9, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(300), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'med': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.6),
                    'max_depth': hyperopt.hp.quniform('max_depth', 4, 7, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(300), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'small': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
                    'max_depth': hyperopt.hp.quniform('max_depth', 3, 6, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(100), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'xsmall': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
                    'max_depth': hyperopt.hp.quniform('max_depth', 3, 6, 1),
                    # Never use for real data, but convenient for tiny sets in test suite
                    'min_child_weight': 0,
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                }
            }[dataset_size]
        })
    ]

    if final_num_trees is not None and final_num_trees != initial_num_trees:
        tune_spaces.append(('trees', {
            'iterations': 30,
            'space': {
                'num_rounds': final_num_trees,
                'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
            }
        }))

    # Baseline parameters to start with. Roughly tuned by what has worked in
    # the past. These vary though depending on number of training samples. These
    # defaults are for the smallest of wikis, which are then overridden for larger
    # wikis
    space = {
        'objective': 'rank:ndcg',
        'eval_metric': 'ndcg@10',
        'num_rounds': initial_num_trees,
        'min_child_weight': 200,
        'max_depth': {
            'xlarge': 7,
            'large': 6,
            'med': 5,
            'small': 4,
            'xsmall': 3,
        }[dataset_size],
        'gamma': 0,
        'subsample': 1.0,
        'colsample_bytree': 0.8,
    }

    tuner = ModelSelection(space, tune_spaces)
    train_func = make_cv_objective(
        partial(train, spark=spark), folds, num_cv_jobs,
        cv_transformer, train_matrix=train_matrix)
    trials_pool = tuner.build_pool(folds, num_cv_jobs)
    return tuner(train_func, trials_pool)

from collections import defaultdict
from functools import partial
import tempfile
from typing import cast, Any, Callable, Dict, List, Mapping, Optional

import hyperopt
import numpy as np
from pyspark.sql import SparkSession
import xgboost as xgb

from mjolnir.training.tuning import make_cv_objective, ModelSelection
from mjolnir.utils import as_local_path, as_local_paths, as_output_file


def _coerce_params(params: Mapping[str, Any]) -> Dict[str, Any]:
    """Force xgboost parameters into appropriate types

    The output from hyperopt is always floats, but some xgboost parameters
    explicitly require integers. Cast those as necessary

    Parameters
    ----------
    params : dict
        xgboost parameters

    Returns
    -------
    dict
        Input parameters coerced as necessary
    """
    def identity(x):
        return x

    def sloppy_int(x):
        try:
            return int(x)
        except ValueError:
            pass

        val = float(x)
        # This could fail for larger numbers due to fp precision, but not
        # expecting integer values larger than two digits here.
        if val.is_integer():
            return int(val)
        raise ValueError('Not parsable as integer: {}'.format(x))

    types = cast(Dict[str, Callable[[Any], Any]], defaultdict(lambda: identity))
    types.update({
        'max_depth': sloppy_int,
        'max_bin': sloppy_int,
        'num_class': sloppy_int,
        'silent': sloppy_int,
    })
    return {k: types[k](v) for k, v in params.items()}


def train(
    fold: Mapping[str, str],
    params: Mapping[str, Any],
    train_matrix: Optional[str] = None,
    spark: Optional[SparkSession] = None
) -> 'XGBoostModel':
    """Train a single xgboost ranking model.

    Primary entry point for hyperparameter tuning normalizes
    parameters and auto detects some values. Actual training is
    passed on to XGBoostModel.trainWithFiles

    Parameters
    ----------
    fold :
        Map from split name to data path. All provided splits will be
        evaluated on each boosting iteration.
    params :
        parameters to pass on to xgboost training
    train_matrix :
        Optional name of training matrix in fold. If not provided will
        auto-detect to either 'all' or 'train'
    spark:
        If provided, train remotely over spark

    Returns
    -------
    XGBoostModel
        Trained xgboost model
    """
    # hyperparameter tuning may have given us floats where we need
    # ints, so this gets all the types right for Java. Also makes
    # a copy of params so we don't modifying the incoming dict.
    # TODO: Does python care about this?  Even if not, in the end it seems
    # reasonable to not pass floats for integer values.
    params = _coerce_params(params)

    # TODO: Maybe num_rounds should just be external? But it's easier
    # to do hyperparameter optimization with a consistent dict interface
    kwargs = cast(Dict[str, Any], {
        'num_boost_round': 100,
    })

    if 'num_boost_round' in params:
        kwargs['num_boost_round'] = params['num_boost_round']
        del params['num_rounds']
    if 'early_stopping_rounds' in params:
        kwargs['early_stopping_rounds'] = params['early_stopping_rounds']
        del params['early_stopping_rounds']

    # Set some sane defaults for ranking tasks
    if 'objective' not in params:
        params['objective'] = 'rank:ndcg'
    if 'eval_metric' not in params:
        params['eval_metric'] = 'ndcg@10'
    # Not really ranking specific, but generally fastest method
    if 'tree_method' not in params:
        params['tree_method'] = 'hist'
    # Convenience for some situations, but typically be explicit
    # about the name of the matrix to train against.
    if train_matrix is None:
        train_matrix = "all" if "all" in fold else "train"

    if spark:
        return XGBoostModel.trainWithFilesRemote(spark, fold, train_matrix, params, **kwargs)
    else:
        return XGBoostModel.trainWithFiles(fold, train_matrix, params, **kwargs)


# Top level: matrix name
# Second level: metric name
# Inner list: stringified per-iteration metric value
EvalsResult = Mapping[str, Mapping[str, List[str]]]


class XGBoostBooster(object):
    """Wrapper for xgb.Booster usage in mjolnir

    Wraps the booster to distinguish what we have after training,
    the XGBoostModel, from what we write to disk, which is only the
    booster. Would be better if there was a clean way to wrap all
    the data up an serialize together, while working with xgboost's
    c++ methods that expect file paths.
    """
    def __init__(self, booster: xgb.Booster) -> None:
        self.booster = booster

    @staticmethod
    def loadBoosterFromHadoopFile(path: str) -> 'XGBoostBooster':
        with as_local_path(path) as local_path:
            return XGBoostBooster.loadBoosterFromLocalFile(local_path)

    @staticmethod
    def loadBoosterFromLocalFile(path: str) -> 'XGBoostBooster':
        booster = xgb.Booster.load_model(path)
        # TODO: Not having the training parameters or the evaluation metrics
        # almost makes this a different thing...
        return XGBoostBooster(booster)

    def saveBoosterAsHadoopFile(self, path: str):
        with as_output_file(path) as f:
            self.saveBoosterAsLocalFile(f.name)

    def saveBoosterAsLocalFile(self, path: str):
        # TODO: This doesn't save any metrics, should it?
        self.booster.save_model(path)


class XGBoostModel(XGBoostBooster):
    """xgboost booster along with train-time metrics

    TODO: Take XGBoostBooster as init arg instead of xgb.Booster?
    """

    def __init__(
        self,
        booster: xgb.Booster,
        evals_result: EvalsResult,
        params: Mapping[str, Any]
    ) -> None:
        super().__init__(booster)
        self.evals_result = evals_result
        self.params = params

    @staticmethod
    def trainWithFilesRemote(
        spark: SparkSession,
        fold: Mapping[str, str],
        train_matrix: str,
        params: Mapping[str, Any],
        **kwargs
    ) -> 'XGBoostModel':
        """Train model on a single remote spark executor.

        Silly hack to train models inside the yarn cluster. To train multiple
        models in parallel python threads will need to be used.  Wish pyspark
        had collectAsync.
        """
        nthread = int(spark.conf.get('spark.task.cpus', '1'))
        if 'nthread' not in params:
            params = dict(params, nthread=nthread)
        elif params['nthread'] != nthread:
            raise Exception("Executors have [{}] cpus but training requested [{}]".format(
                nthread, params['nthread']))

        return (
            spark.sparkContext
            .parallelize([1], 1)
            .map(lambda x: XGBoostModel.trainWithFiles(fold, train_matrix, params, **kwargs))
            .collect()[0]
        )

    @staticmethod
    def trainWithFiles(
        fold: Mapping[str, str],
        train_matrix: str,
        params: Mapping[str, Any],
        **kwargs
    ) -> 'XGBoostModel':
        """Wrapper around xgb.train

        This intentionally forwards to trainWithRDD, rather than
        trainWithDataFrame, as the underlying method currently prevents using
        rank:pairwise and metrics with @, such as ndcg@5.

        Parameters
        ----------
        fold :
            Map from split name to data path. All provided splits will be
            evaluated on each boosting iteration.
        train_matrix: str
            name of split in fold to train against
        params : dict
            XGBoost training parameters

        Returns
        -------
        mjolnir.training.xgboost.XGBoostModel
            trained xgboost ranking model
        """
        with as_local_paths(fold.values()) as local_paths:
            matrices = {name: xgb.DMatrix(path) for name, path in zip(fold.keys(), local_paths)}
            dtrain = matrices[train_matrix]
            evallist = [(dmat, name) for name, dmat in matrices.items()]
            metrics = cast(Mapping, {})
            booster = xgb.train(params, dtrain, evals=evallist, evals_result=metrics, **kwargs)
            return XGBoostModel(booster, metrics, params)

    def dump(self, features=None, with_stats=False, format="json"):
        """Dumps the xgboost model

        Parameters
        ----------
        features : list of str or None, optional
            list of features names, or None for no feature names in dump.
            (Default: None)
        withStats : bool, optional
            Should various additional statistics be included? These are not
            necessary for prediction. (Default: False)
        format : string, optional
            The format of dump to produce, either json or text. (Default: json)

        Returns
        -------
        str
            valid json string containing all trees
        """
        # Annoyingly the xgboost api doesn't take the feature map as a string, but
        # instead as a filename. Write the feature map out to a file if necessary.
        if features:
            # When we write the svmrank formatted files the features are indexed
            # starting at 1. We need to throw a fake index 0 in here or it's all
            # wrong.
            feat_map = "0 PLACEHOLDER_FEAT q\n" + \
                "\n".join(["%d %s q" % (i, fname) for i, fname in enumerate(features, 1)])
            fmap_f = tempfile.NamedTemporaryFile(mode='w')
            fmap_f.write(feat_map)
            fmap_f.flush()
            fmap_path = fmap_f.name
        else:
            fmap_path = ''

        trees = self.booster.get_dump(fmap_path, with_stats, dump_format='json')
        # For whatever reason we get a json line per tree. Turn that into an array
        # so we have a single valid json string.
        return '[' + ','.join(trees) + ']'


def cv_transformer(model: XGBoostModel, params: Mapping[str, Any]):
    """Report model metrics in format expected by model selection"""
    metric = params['eval_metric']
    return {
        'train': model.evals_result['train'][metric][-1],
        'test': model.evals_result['test'][metric][-1],
        'metrics': model.evals_result,
    }


def tune(
    folds: List[Mapping[str, str]],
    stats: Dict,
    train_matrix: str,
    num_cv_jobs: int = 5,
    initial_num_trees: int = 100,
    final_num_trees: int = 500,
    iterations: int = 150,
    spark: Optional[SparkSession] = None
):
    """Find appropriate hyperparameters for training df

    This is far from perfect, hyperparameter tuning is a bit of a black art
    and could probably benefit from human interaction at each stage. Various
    parameters depend a good bit on the number of samples in df, and how
    that data is shaped.

    Below is tuned for a dataframe with approximatly 10k normalized queries,
    110k total queries, and 2.2M samples. This is actually a relatively small
    dataset, we should rework the values used with larger data sets if they
    are promising. It may also be that the current feature space can't take
    advantage of more samples.

    Note that hyperopt uses the first 20 iterations to initialize, during those
    first 20 this is a strictly random search.

    Parameters
    ----------
    folds : list of dict containing train and test keys
    stats : dict
        stats about the fold from the make_folds utility script
    num_cv_jobs : int, optional
        The number of cross validation folds to train in parallel. (Default: 5)
    initial_num_trees: int, optional
        The number of trees to do most of the hyperparameter tuning with. This
        should be large enough to be resonably representative of the final
        training size. (Default: 100)
    final_num_trees: int, optional
        The number of trees to do the final eta optimization with. If set to
        None the final eta optimization will be skipped and initial_n_tree will
        be kept.

    Returns
    -------
    dict
        Dict with two keys, trials and params. params is the optimal set of
        parameters. trials contains a dict of individual optimization steps
        performed, each containing a hyperopt.Trials object recording what
        happened.
    """
    num_obs = stats['num_observations']

    if num_obs > 8000000:
        dataset_size = 'xlarge'
    elif num_obs > 1000000:
        dataset_size = 'large'
    elif num_obs > 500000:
        dataset_size = 'med'
    elif num_obs > 500:
        dataset_size = 'small'
    else:
        dataset_size = 'xsmall'

    # Setup different tuning profiles for different sizes of datasets.
    tune_spaces = [
        ('initial', {
            'iterations': iterations,
            'space': {
                'xlarge': {
                    'eta': hyperopt.hp.uniform('eta', 0.3, 0.8),
                    # Have seen values of 7 and 10 as best on roughly same size
                    # datasets from different wikis. It really just depends.
                    'max_depth': hyperopt.hp.quniform('max_depth', 6, 11, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(500), 10),
                    # % of features to use for each tree. helps prevent overfit
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'large': {
                    'eta': hyperopt.hp.uniform('eta', 0.3, 0.6),
                    'max_depth': hyperopt.hp.quniform('max_depth', 5, 9, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(300), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'med': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.6),
                    'max_depth': hyperopt.hp.quniform('max_depth', 4, 7, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(300), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'small': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
                    'max_depth': hyperopt.hp.quniform('max_depth', 3, 6, 1),
                    'min_child_weight': hyperopt.hp.qloguniform(
                        'min_child_weight', np.log(10), np.log(100), 10),
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                },
                'xsmall': {
                    'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
                    'max_depth': hyperopt.hp.quniform('max_depth', 3, 6, 1),
                    # Never use for real data, but convenient for tiny sets in test suite
                    'min_child_weight': 0,
                    'colsample_bytree': hyperopt.hp.quniform('colsample_bytree', 0.8, 1, .01),
                    'subsample': hyperopt.hp.quniform('subsample', 0.8, 1, .01),
                }
            }[dataset_size]
        })
    ]

    if final_num_trees is not None and final_num_trees != initial_num_trees:
        tune_spaces.append(('trees', {
            'iterations': 30,
            'space': {
                'num_rounds': final_num_trees,
                'eta': hyperopt.hp.uniform('eta', 0.1, 0.4),
            }
        }))

    # Baseline parameters to start with. Roughly tuned by what has worked in
    # the past. These vary though depending on number of training samples. These
    # defaults are for the smallest of wikis, which are then overridden for larger
    # wikis
    space = {
        'objective': 'rank:ndcg',
        'eval_metric': 'ndcg@10',
        'num_rounds': initial_num_trees,
        'min_child_weight': 200,
        'max_depth': {
            'xlarge': 7,
            'large': 6,
            'med': 5,
            'small': 4,
            'xsmall': 3,
        }[dataset_size],
        'gamma': 0,
        'subsample': 1.0,
        'colsample_bytree': 0.8,
    }

    tuner = ModelSelection(space, tune_spaces)
    train_func = make_cv_objective(
        partial(train, spark=spark), folds, num_cv_jobs,
        cv_transformer, train_matrix=train_matrix)
    trials_pool = tuner.build_pool(folds, num_cv_jobs)
    return tuner(train_func, trials_pool)
Force xgboost parameters into appropriate types      The output from hyperopt is always floats, but some xgboost parameters     explicitly require integers. Cast those as necessary      Parameters     ----------     params : dict         xgboost parameters      Returns     -------     dict         Input parameters coerced as necessary # This could fail for larger numbers due to fp precision, but not # expecting integer values larger than two digits here. Train a single xgboost ranking model.      Primary entry point for hyperparameter tuning normalizes     parameters and auto detects some values. Actual training is     passed on to XGBoostModel.trainWithFiles      Parameters     ----------     fold :         Map from split name to data path. All provided splits will be         evaluated on each boosting iteration.     params :         parameters to pass on to xgboost training     train_matrix :         Optional name of training matrix in fold. If not provided will         auto-detect to either 'all' or 'train'     spark:         If provided, train remotely over spark      Returns     -------     XGBoostModel         Trained xgboost model # hyperparameter tuning may have given us floats where we need # ints, so this gets all the types right for Java. Also makes # a copy of params so we don't modifying the incoming dict. # TODO: Does python care about this?  Even if not, in the end it seems # reasonable to not pass floats for integer values. # TODO: Maybe num_rounds should just be external? But it's easier # to do hyperparameter optimization with a consistent dict interface # Set some sane defaults for ranking tasks # Not really ranking specific, but generally fastest method # Convenience for some situations, but typically be explicit # about the name of the matrix to train against. # Top level: matrix name # Second level: metric name # Inner list: stringified per-iteration metric value Wrapper for xgb.Booster usage in mjolnir      Wraps the booster to distinguish what we have after training,     the XGBoostModel, from what we write to disk, which is only the     booster. Would be better if there was a clean way to wrap all     the data up an serialize together, while working with xgboost's     c++ methods that expect file paths. # TODO: Not having the training parameters or the evaluation metrics # almost makes this a different thing... # TODO: This doesn't save any metrics, should it? xgboost booster along with train-time metrics      TODO: Take XGBoostBooster as init arg instead of xgb.Booster? Train model on a single remote spark executor.          Silly hack to train models inside the yarn cluster. To train multiple         models in parallel python threads will need to be used.  Wish pyspark         had collectAsync. Wrapper around xgb.train          This intentionally forwards to trainWithRDD, rather than         trainWithDataFrame, as the underlying method currently prevents using         rank:pairwise and metrics with @, such as ndcg@5.          Parameters         ----------         fold :             Map from split name to data path. All provided splits will be             evaluated on each boosting iteration.         train_matrix: str             name of split in fold to train against         params : dict             XGBoost training parameters          Returns         -------         mjolnir.training.xgboost.XGBoostModel             trained xgboost ranking model Dumps the xgboost model          Parameters         ----------         features : list of str or None, optional             list of features names, or None for no feature names in dump.             (Default: None)         withStats : bool, optional             Should various additional statistics be included? These are not             necessary for prediction. (Default: False)         format : string, optional             The format of dump to produce, either json or text. (Default: json)          Returns         -------         str             valid json string containing all trees # Annoyingly the xgboost api doesn't take the feature map as a string, but # instead as a filename. Write the feature map out to a file if necessary. # When we write the svmrank formatted files the features are indexed # starting at 1. We need to throw a fake index 0 in here or it's all # wrong. # For whatever reason we get a json line per tree. Turn that into an array # so we have a single valid json string. Report model metrics in format expected by model selection Find appropriate hyperparameters for training df      This is far from perfect, hyperparameter tuning is a bit of a black art     and could probably benefit from human interaction at each stage. Various     parameters depend a good bit on the number of samples in df, and how     that data is shaped.      Below is tuned for a dataframe with approximatly 10k normalized queries,     110k total queries, and 2.2M samples. This is actually a relatively small     dataset, we should rework the values used with larger data sets if they     are promising. It may also be that the current feature space can't take     advantage of more samples.      Note that hyperopt uses the first 20 iterations to initialize, during those     first 20 this is a strictly random search.      Parameters     ----------     folds : list of dict containing train and test keys     stats : dict         stats about the fold from the make_folds utility script     num_cv_jobs : int, optional         The number of cross validation folds to train in parallel. (Default: 5)     initial_num_trees: int, optional         The number of trees to do most of the hyperparameter tuning with. This         should be large enough to be resonably representative of the final         training size. (Default: 100)     final_num_trees: int, optional         The number of trees to do the final eta optimization with. If set to         None the final eta optimization will be skipped and initial_n_tree will         be kept.      Returns     -------     dict         Dict with two keys, trials and params. params is the optimal set of         parameters. trials contains a dict of individual optimization steps         performed, each containing a hyperopt.Trials object recording what         happened. # Setup different tuning profiles for different sizes of datasets. # Have seen values of 7 and 10 as best on roughly same size # datasets from different wikis. It really just depends. # % of features to use for each tree. helps prevent overfit # Never use for real data, but convenient for tiny sets in test suite # Baseline parameters to start with. Roughly tuned by what has worked in # the past. These vary though depending on number of training samples. These # defaults are for the smallest of wikis, which are then overridden for larger # wikis