<commit_before>import os
import sys
import logging
import inspect

import pandas as pd
import numpy as np

from sklearn.externals import joblib
from sklearn.base import _pprint

from drain import util, metrics
from drain.step import Step, Construct


class FitPredict(Step):
    """
    Step which can fit a scikit-learn estimator and make predictions.
    """
    def __init__(self, inputs,
                 return_estimator=False,
                 return_feature_importances=True,
                 return_predictions=True,
                 prefit=False,
                 predict_train=False):
        """
        Args:
            return_estimator: whether to return the fitted estimator object
            return_feature_importances: whether to return a DataFrame of feature importances
            prefit: whether the estimator input is already fitted
            predict_train: whether to make predictions on training set
        """
        Step.__init__(self, inputs=inputs, return_estimator=return_estimator,
                      return_feature_importances=return_feature_importances,
                      return_predictions=return_predictions, prefit=prefit,
                      predict_train=predict_train)

    def run(self, estimator, X, y, train=None, test=None, aux=None, sample_weight=None,
            feature_importances=None):
        if not self.prefit:
            if train is not None:
                X_train, y_train = X[train], y[train]
            else:
                X_train, y_train = X, y

            y_train = y_train.astype(bool)

            logging.info('Fitting with %s examples, %s features' % X_train.shape)
            if 'sample_weight' in inspect.getargspec(estimator.fit).args and\
                    sample_weight is not None:
                logging.info('Using sample weight')
                sample_weight = sample_weight.ix[train.index][train].values
                estimator.fit(X_train, y_train, sample_weight=sample_weight)
            else:
                estimator.fit(X_train, y_train)

        result = {}

        if self.return_estimator:
            result['estimator'] = estimator
        if self.return_feature_importances:
            result['feature_importances'] = feature_importance(estimator, X)
        if self.return_predictions:
            if test is not None and not self.predict_train:
                X_test, y_test = X[test], y[test]
            else:
                X_test, y_test = X, y

            logging.info('Predicting %s examples' % len(X_test))
            y = pd.DataFrame({'true': y_test})
            y['score'] = y_score(estimator, X_test)

            if self.predict_train:
                y['train'] = train

            if aux is not None:
                y = y.join(aux, how='left')

            result['y'] = y

        return result

    def dump(self):
        result = self.get_result()
        if self.return_estimator:
            filename = os.path.join(self._dump_dirname, 'estimator.pkl')
            joblib.dump(result['estimator'], filename)
        if self.return_feature_importances:
            filename = os.path.join(self._dump_dirname, 'feature_importances.hdf')
            result['feature_importances'].to_hdf(filename, 'df')
        if self.return_predictions:
            filename = os.path.join(self._dump_dirname, 'y.hdf')
            result['y'].to_hdf(filename, 'df')

    def load(self):
        result = {}
        if self.return_estimator:
            filename = os.path.join(self._dump_dirname, 'estimator.pkl')
            result['estimator'] = joblib.load(filename)
        if self.return_feature_importances:
            filename = os.path.join(self._dump_dirname, 'feature_importances.hdf')
            result['feature_importances'] = pd.read_hdf(filename, 'df')
        if self.return_predictions:
            filename = os.path.join(self._dump_dirname, 'y.hdf')
            result['y'] = pd.read_hdf(filename, 'df')

        self.set_result(result)


class Fit(FitPredict):
    def __init__(self, inputs, return_estimator=True, return_feature_importances=False):
        FitPredict.__init__(self, inputs=inputs, prefit=False,
                            return_estimator=return_estimator,
                            return_feature_importances=return_feature_importances,
                            return_predictions=False)


class Predict(FitPredict):
    def __init__(self, inputs, return_estimator=False, return_feature_importances=False):
        FitPredict.__init__(self, inputs=inputs,
                            return_feature_importances=return_feature_importances,
                            return_estimator=return_estimator,
                            return_predictions=True, prefit=True)


class PredictProduct(Step):
    def run(self, **kwargs):
        keys = kwargs.keys()
        ys = [kwargs[k]['y'] for k in keys]
        y = ys[0].copy()
        y.rename(columns={'score': 'score_%s' % keys[0]}, inplace=True)
        y['score_%s' % keys[1]] = ys[1].score
        y['score'] = ys[0].score * ys[1].score

        return {'y': y}


class InverseProbabilityWeights(Step):
    def run(self, y, train=None, **kwargs):
        if train is not None:
            logging.info("Using training mask")
            train = train[train].index
            intersection = y.index.intersection(train)
            if len(intersection) != len(train):
                raise ValueError("Must provide scores for every training example.")
            y = y.ix[intersection]

        return {'sample_weight': y.score**-1}


def y_score(estimator, X):
    """
    Score examples from a new matrix X
    Args:
        estimator: an sklearn estimator object
        X: design matrix with the same features that the estimator was trained on

    Returns: a vector of scores of the same length as X
<commit_msg>keys list<commit_after>import os
import sys
import logging
import inspect

import pandas as pd
import numpy as np

from sklearn.externals import joblib
from sklearn.base import _pprint

from drain import util, metrics
from drain.step import Step, Construct


class FitPredict(Step):
    """
    Step which can fit a scikit-learn estimator and make predictions.
    """
    def __init__(self, inputs,
                 return_estimator=False,
                 return_feature_importances=True,
                 return_predictions=True,
                 prefit=False,
                 predict_train=False):
        """
        Args:
            return_estimator: whether to return the fitted estimator object
            return_feature_importances: whether to return a DataFrame of feature importances
            prefit: whether the estimator input is already fitted
            predict_train: whether to make predictions on training set
        """
        Step.__init__(self, inputs=inputs, return_estimator=return_estimator,
                      return_feature_importances=return_feature_importances,
                      return_predictions=return_predictions, prefit=prefit,
                      predict_train=predict_train)

    def run(self, estimator, X, y, train=None, test=None, aux=None, sample_weight=None,
            feature_importances=None):
        if not self.prefit:
            if train is not None:
                X_train, y_train = X[train], y[train]
            else:
                X_train, y_train = X, y

            y_train = y_train.astype(bool)

            logging.info('Fitting with %s examples, %s features' % X_train.shape)
            if 'sample_weight' in inspect.getargspec(estimator.fit).args and\
                    sample_weight is not None:
                logging.info('Using sample weight')
                sample_weight = sample_weight.ix[train.index][train].values
                estimator.fit(X_train, y_train, sample_weight=sample_weight)
            else:
                estimator.fit(X_train, y_train)

        result = {}

        if self.return_estimator:
            result['estimator'] = estimator
        if self.return_feature_importances:
            result['feature_importances'] = feature_importance(estimator, X)
        if self.return_predictions:
            if test is not None and not self.predict_train:
                X_test, y_test = X[test], y[test]
            else:
                X_test, y_test = X, y

            logging.info('Predicting %s examples' % len(X_test))
            y = pd.DataFrame({'true': y_test})
            y['score'] = y_score(estimator, X_test)

            if self.predict_train:
                y['train'] = train

            if aux is not None:
                y = y.join(aux, how='left')

            result['y'] = y

        return result

    def dump(self):
        result = self.get_result()
        if self.return_estimator:
            filename = os.path.join(self._dump_dirname, 'estimator.pkl')
            joblib.dump(result['estimator'], filename)
        if self.return_feature_importances:
            filename = os.path.join(self._dump_dirname, 'feature_importances.hdf')
            result['feature_importances'].to_hdf(filename, 'df')
        if self.return_predictions:
            filename = os.path.join(self._dump_dirname, 'y.hdf')
            result['y'].to_hdf(filename, 'df')

    def load(self):
        result = {}
        if self.return_estimator:
            filename = os.path.join(self._dump_dirname, 'estimator.pkl')
            result['estimator'] = joblib.load(filename)
        if self.return_feature_importances:
            filename = os.path.join(self._dump_dirname, 'feature_importances.hdf')
            result['feature_importances'] = pd.read_hdf(filename, 'df')
        if self.return_predictions:
            filename = os.path.join(self._dump_dirname, 'y.hdf')
            result['y'] = pd.read_hdf(filename, 'df')

        self.set_result(result)


class Fit(FitPredict):
    def __init__(self, inputs, return_estimator=True, return_feature_importances=False):
        FitPredict.__init__(self, inputs=inputs, prefit=False,
                            return_estimator=return_estimator,
                            return_feature_importances=return_feature_importances,
                            return_predictions=False)


class Predict(FitPredict):
    def __init__(self, inputs, return_estimator=False, return_feature_importances=False):
        FitPredict.__init__(self, inputs=inputs,
                            return_feature_importances=return_feature_importances,
                            return_estimator=return_estimator,
                            return_predictions=True, prefit=True)


class PredictProduct(Step):
    def run(self, **kwargs):
        keys = list(kwargs.keys())
        ys = [kwargs[k]['y'] for k in keys]
        y = ys[0].copy()
        y.rename(columns={'score': 'score_%s' % keys[0]}, inplace=True)
        y['score_%s' % keys[1]] = ys[1].score
        y['score'] = ys[0].score * ys[1].score

        return {'y': y}


class InverseProbabilityWeights(Step):
    def run(self, y, train=None, **kwargs):
        if train is not None:
            logging.info("Using training mask")
            train = train[train].index
            intersection = y.index.intersection(train)
            if len(intersection) != len(train):
                raise ValueError("Must provide scores for every training example.")
            y = y.ix[intersection]

        return {'sample_weight': y.score**-1}


def y_score(estimator, X):
    """
    Score examples from a new matrix X
    Args:
        estimator: an sklearn estimator object
        X: design matrix with the same features that the estimator was trained on

    Returns: a vector of scores of the same length as X

{'fraction_non_alphanumeric': 0.09412972787951394, 'fraction_numerical': 0.001198014718466541, 'mean_word_length': 2.1862050163576883, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1402559', 'n_tokens_mistral': 3356, 'n_tokens_neox': 3258, 'n_words': 775}