jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def fit_baseline_model(self, train_data: Optional[pd.DataFrame] = None, train_labels: Optional[pd.Series] = None) -> BaseEstimator:
        """
        Fit a baseline model. If no data is given (default), it uses the task's train data and creates the attribute `_baseline_model`. \
            If data is given, it trains this data.

        Args:
            train_data (Optional[pd.DataFrame], optional): Data to train. Defaults to None.
            train_labels (Optional[pd.Series], optional): Labels to train. Defaults to None.

        Raises:
            ValueError: If `train_data` is given but `train_labels` not or vice versa

        Returns:
            BaseEstimator: Trained model
        """

        if (train_data is None and train_labels is not None) or (train_data is not None and train_labels is None):
            raise ValueError("either set both parameters (train_data, train_labels) or non")

        use_original_data = train_data is None

        # shortcut if model is already trained
        if use_original_data and self._baseline_model:
            return self._baseline_model

        if use_original_data:
            train_data = self.train_data.copy()
            train_labels = self.train_labels.copy()

        for col in self.categorical_columns:
            train_data[col] = train_data[col].astype(str)

        categorical_preprocessing = Pipeline(
            [
                ('mark_missing', Imputer(strategy='most_frequent')),
                ('one_hot_encode', OneHotEncoder(handle_unknown='ignore'))
            ]
        )

        numerical_preprocessing = Pipeline(
            [
                ('mark_missing', Imputer(strategy='mean')),
                ('scaling',  StandardScaler())
            ]
        )

        feature_transformation = ColumnTransformer(transformers=[
                ('categorical_features', categorical_preprocessing, self.categorical_columns),
                ('scaled_numeric', numerical_preprocessing, self.numerical_columns)
            ]
        )

        param_grid, pipeline, scorer = self._get_pipeline_grid_scorer_tuple(feature_transformation)
        refit = list(scorer.keys())[0]

        search = GridSearchCV(pipeline, param_grid, scoring=scorer, n_jobs=-1, refit=refit)
        model = search.fit(train_data, train_labels).best_estimator_

        # only set baseline model attribute if it is trained on the original task data
        if use_original_data:
            self._baseline_model = model

        return model
//  [python]
def fit_baseline_model(self, train_data: Optional[pd.DataFrame] = None, train_labels: Optional[pd.Series] = None) -> BaseEstimator:
        """
        Fit a baseline model. If no data is given (default), it uses the task's train data and creates the attribute `_baseline_model`. \
            If data is given, it trains this data.

        Args:
            train_data (Optional[pd.DataFrame], optional): Data to train. Defaults to None.
            train_labels (Optional[pd.Series], optional): Labels to train. Defaults to None.

        Raises:
            ValueError: If `train_data` is given but `train_labels` not or vice versa

        Returns:
            BaseEstimator: Trained model
        """

        if (train_data is None and train_labels is not None) or (train_data is not None and train_labels is None):
            raise ValueError("either set both parameters (train_data, train_labels) or non")

        use_original_data = train_data is None

        # shortcut if model is already trained
        if use_original_data and self._baseline_model:
            return self._baseline_model

        if use_original_data:
            train_data = self.train_data.copy()
            train_labels = self.train_labels.copy()

        for col in self.categorical_columns:
            train_data[col] = train_data[col].astype(str)

        categorical_preprocessing = Pipeline(
            [
                ('mark_missing', SimpleImputer(strategy='most_frequent')),
                ('one_hot_encode', OneHotEncoder(handle_unknown='ignore'))
            ]
        )

        numerical_preprocessing = Pipeline(
            [
                ('mark_missing', SimpleImputer(strategy='mean')),
                ('scaling',  StandardScaler())
            ]
        )

        feature_transformation = ColumnTransformer(transformers=[
                ('categorical_features', categorical_preprocessing, self.categorical_columns),
                ('scaled_numeric', numerical_preprocessing, self.numerical_columns)
            ]
        )

        param_grid, pipeline, scorer = self._get_pipeline_grid_scorer_tuple(feature_transformation)
        refit = list(scorer.keys())[0]

        search = GridSearchCV(pipeline, param_grid, scoring=scorer, n_jobs=-1, refit=refit)
        model = search.fit(train_data, train_labels).best_estimator_

        # only set baseline model attribute if it is trained on the original task data
        if use_original_data:
            self._baseline_model = model

        return model
//  [python]
def clustered_map(n_clicks, year, n_clusters, indicators):
    if not indicators:
        raise PreventUpdate
    imp = Imputer(missing_values=np.nan, strategy='mean')
    scaler = StandardScaler()
    kmeans = KMeans(n_clusters=n_clusters)
    
    df = poverty[poverty['is_country'] & poverty['year'].eq(year)][indicators + ['Country Name', 'year']]
    data = df[indicators]
    if df.isna().all().any():
        return px.scatter(title='No available data for the selected combination of year/indicators.')
    data_no_na = imp.fit_transform(data)
    scaled_data = scaler.fit_transform(data_no_na)
    kmeans.fit(scaled_data)

    fig = px.choropleth(df,
                        locations='Country Name',
                        locationmode='country names',
                        color=[str(x) for x in  kmeans.labels_], 
                        labels={'color': 'Cluster'},
                        hover_data=indicators,
                        height=650,
                        title=f'Country clusters - {year}. Number of clusters: {n_clusters}<br>Inertia: {kmeans.inertia_:,.2f}',
                        color_discrete_sequence=px.colors.qualitative.T10)
    fig.add_annotation(x=-0.1, y=-0.15, 
                       xref='paper', yref='paper',
                       text='Indicators:<br>' + "<br>".join(indicators), 
                       showarrow=False)
    fig.layout.geo.showframe = False
    fig.layout.geo.showcountries = True
    fig.layout.geo.projection.type = 'natural earth'
    fig.layout.geo.lataxis.range = [-53, 76]
    fig.layout.geo.lonaxis.range = [-137, 168]
    fig.layout.geo.landcolor = 'white'
    fig.layout.geo.bgcolor = '#E5ECF6'
    fig.layout.paper_bgcolor = '#E5ECF6'
    fig.layout.geo.countrycolor = 'gray'
    fig.layout.geo.coastlinecolor = 'gray'
    return fig
//  [python]
def clustered_map(n_clicks, year, n_clusters, indicators):
    if not indicators:
        raise PreventUpdate
    imp = SimpleImputer(missing_values=np.nan, strategy='mean')
    scaler = StandardScaler()
    kmeans = KMeans(n_clusters=n_clusters)
    
    df = poverty[poverty['is_country'] & poverty['year'].eq(year)][indicators + ['Country Name', 'year']]
    data = df[indicators]
    if df.isna().all().any():
        return px.scatter(title='No available data for the selected combination of year/indicators.')
    data_no_na = imp.fit_transform(data)
    scaled_data = scaler.fit_transform(data_no_na)
    kmeans.fit(scaled_data)

    fig = px.choropleth(df,
                        locations='Country Name',
                        locationmode='country names',
                        color=[str(x) for x in  kmeans.labels_], 
                        labels={'color': 'Cluster'},
                        hover_data=indicators,
                        height=650,
                        title=f'Country clusters - {year}. Number of clusters: {n_clusters}<br>Inertia: {kmeans.inertia_:,.2f}',
                        color_discrete_sequence=px.colors.qualitative.T10)
    fig.add_annotation(x=-0.1, y=-0.15, 
                       xref='paper', yref='paper',
                       text='Indicators:<br>' + "<br>".join(indicators), 
                       showarrow=False)
    fig.layout.geo.showframe = False
    fig.layout.geo.showcountries = True
    fig.layout.geo.projection.type = 'natural earth'
    fig.layout.geo.lataxis.range = [-53, 76]
    fig.layout.geo.lonaxis.range = [-137, 168]
    fig.layout.geo.landcolor = 'white'
    fig.layout.geo.bgcolor = '#E5ECF6'
    fig.layout.paper_bgcolor = '#E5ECF6'
    fig.layout.geo.countrycolor = 'gray'
    fig.layout.geo.coastlinecolor = 'gray'
    return fig
//  [python]
def log_artifacts_and_results() -> (
    tuple[np.ndarray, pd.DataFrame, str, dict, list, int, str, Pipeline]
):
    encoder_to_imputer = Pipeline(
        steps=[
            (
                "imputer",
                Imputer(missing_values="", strategy="constant", fill_value="C"),
            ),
            ("encoder", OrdinalEncoder()),
        ]
    )
    encoder_to_imputer.fit([["A"], ["B"], ["C"]])

    context = mlrun.get_or_create_ctx(name="ctx")
    context.log_result(key="manually_logged_result", value=10)

    file_path = os.path.join(context.artifact_path, "my_file.txt")
    with open(file_path, "w") as file:
        file.write("123")

    return (
        np.ones((10, 20)),
        pd.DataFrame(np.zeros((20, 10))),
        file_path,
        {"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]},
        [["A"], ["B"], [""]],
        3,
        "hello",
        encoder_to_imputer,
    )
//  [python]
def log_artifacts_and_results() -> (
    tuple[np.ndarray, pd.DataFrame, str, dict, list, int, str, Pipeline]
):
    encoder_to_imputer = Pipeline(
        steps=[
            (
                "imputer",
                SimpleImputer(missing_values="", strategy="constant", fill_value="C"),
            ),
            ("encoder", OrdinalEncoder()),
        ]
    )
    encoder_to_imputer.fit([["A"], ["B"], ["C"]])

    context = mlrun.get_or_create_ctx(name="ctx")
    context.log_result(key="manually_logged_result", value=10)

    file_path = os.path.join(context.artifact_path, "my_file.txt")
    with open(file_path, "w") as file:
        file.write("123")

    return (
        np.ones((10, 20)),
        pd.DataFrame(np.zeros((20, 10))),
        file_path,
        {"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]},
        [["A"], ["B"], [""]],
        3,
        "hello",
        encoder_to_imputer,
    )
//  [python]
num_pipe = make_pipeline(
    StandardScaler(), Imputer(strategy="mean", add_indicator=True)
)
cat_pipe = make_pipeline(
    Imputer(strategy="constant", fill_value="missing"),
    OneHotEncoder(handle_unknown="ignore"),
)
//  [python]
num_pipe = make_pipeline(
    StandardScaler(), SimpleImputer(strategy="mean", add_indicator=True)
)
cat_pipe = make_pipeline(
    SimpleImputer(strategy="constant", fill_value="missing"),
    OneHotEncoder(handle_unknown="ignore"),
)
//  [python]
def load_mimic(base_dir='./data'):
    data = pd.read_csv(f'{base_dir}/mimic-ii/full_cohort_data.csv')
    # data.drop('hgb_first')
    fs = [
        'aline_flg',
        'gender_num',
        # 'hosp_exp_flg',
        # 'icu_exp_flg',
        # 'day_28_flg',
        # 'censor_flg',
        'sepsis_flg', 'chf_flg', 'afib_flg',
        'renal_flg', 'liver_flg', 'copd_flg', 'cad_flg', 'stroke_flg',
        'mal_flg', 'resp_flg',
    ]
    features_name = fs
    data1 = data[fs].values
    imp_mean = Imputer(missing_values=np.nan, strategy='mean')
    data1 = imp_mean.fit_transform(data1)

    f2 = fs.copy()
    f2.append('day_icu_intime')
    f2.append('service_unit')
    f2.append('day_28_flg')
    f2.append('hospital_los_day')
    f2.append('icu_exp_flg')
    f2.append('hosp_exp_flg')
    f2.append('censor_flg')
    f2.append('mort_day_censored')
    f2 = data.columns.difference(f2)
    data2 = data[f2].values
    imp_mean = Imputer(missing_values=np.nan, strategy='mean')
    data2 = imp_mean.fit_transform(data2)
    scaler = MinMaxScaler((0, 1))
    data2 = scaler.fit_transform(data2)
    features_name = features_name + list(f2)
    est = KBinsDiscretizer(n_bins=3, encode='onehot-dense', strategy='uniform')
    data2d = est.fit_transform(data2)
    f2d = []
    for feature in f2:
        # f2d.append(feature + '_VLOW')
        f2d.append(feature + '_LOW')
        f2d.append(feature + '_NORMAL')
        f2d.append(feature + '_HIGH')
        # f2d.append(feature + '_VHIGH')
    features_name = fs + f2d

    datax = np.hstack((data1, data2d))
    # datay = data['day_28_flg'].values
    # datay = (data['hospital_los_day']>6).values
    # datay = data['hosp_exp_flg'].values

    datay = (data['day_28_flg'].values + data['hosp_exp_flg'].values + data['icu_exp_flg'].values + (1-data['censor_flg'].values)) > 0

    # model = DecisionTreeClassifier()
    # # model = RandomForestClassifier()
    # scores = cross_val_score(model, datax, datay, cv=10)
    # print(scores.mean())

    x = torch.FloatTensor(datax)
    y = one_hot(torch.tensor(datay).to(torch.long)).to(torch.float)
    class_names = ['death_before_28_days', 'survived_28_days']
    return x, y, features_name, class_names
//  [python]
def load_mimic(base_dir='./data'):
    data = pd.read_csv(f'{base_dir}/mimic-ii/full_cohort_data.csv')
    # data.drop('hgb_first')
    fs = [
        'aline_flg',
        'gender_num',
        # 'hosp_exp_flg',
        # 'icu_exp_flg',
        # 'day_28_flg',
        # 'censor_flg',
        'sepsis_flg', 'chf_flg', 'afib_flg',
        'renal_flg', 'liver_flg', 'copd_flg', 'cad_flg', 'stroke_flg',
        'mal_flg', 'resp_flg',
    ]
    features_name = fs
    data1 = data[fs].values
    imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
    data1 = imp_mean.fit_transform(data1)

    f2 = fs.copy()
    f2.append('day_icu_intime')
    f2.append('service_unit')
    f2.append('day_28_flg')
    f2.append('hospital_los_day')
    f2.append('icu_exp_flg')
    f2.append('hosp_exp_flg')
    f2.append('censor_flg')
    f2.append('mort_day_censored')
    f2 = data.columns.difference(f2)
    data2 = data[f2].values
    imp_mean = SimpleImputer(missing_values=np.nan, strategy='mean')
    data2 = imp_mean.fit_transform(data2)
    scaler = MinMaxScaler((0, 1))
    data2 = scaler.fit_transform(data2)
    features_name = features_name + list(f2)
    est = KBinsDiscretizer(n_bins=3, encode='onehot-dense', strategy='uniform')
    data2d = est.fit_transform(data2)
    f2d = []
    for feature in f2:
        # f2d.append(feature + '_VLOW')
        f2d.append(feature + '_LOW')
        f2d.append(feature + '_NORMAL')
        f2d.append(feature + '_HIGH')
        # f2d.append(feature + '_VHIGH')
    features_name = fs + f2d

    datax = np.hstack((data1, data2d))
    # datay = data['day_28_flg'].values
    # datay = (data['hospital_los_day']>6).values
    # datay = data['hosp_exp_flg'].values

    datay = (data['day_28_flg'].values + data['hosp_exp_flg'].values + data['icu_exp_flg'].values + (1-data['censor_flg'].values)) > 0

    # model = DecisionTreeClassifier()
    # # model = RandomForestClassifier()
    # scores = cross_val_score(model, datax, datay, cv=10)
    # print(scores.mean())

    x = torch.FloatTensor(datax)
    y = one_hot(torch.tensor(datay).to(torch.long)).to(torch.float)
    class_names = ['death_before_28_days', 'survived_28_days']
    return x, y, features_name, class_names
//  [python]
def setUp(self) -> None:
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

        with open(os.path.join(ROOT_DIR, '..', 'dtypes.json'), 'r') as jsonfile:
            self.dtyp = json.load(jsonfile)
        self.data = pd.read_csv(os.path.join(ROOT_DIR, '..', 'Datasets', 'earthquake_data.csv'),
                                dtype = self.dtyp)
        self.imp = Imputer(missing_values = np.nan, strategy = 'constant', fill_value = 'NA')
//  [python]
def setUp(self) -> None:
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

        with open(os.path.join(ROOT_DIR, '..', 'dtypes.json'), 'r') as jsonfile:
            self.dtyp = json.load(jsonfile)
        self.data = pd.read_csv(os.path.join(ROOT_DIR, '..', 'Datasets', 'earthquake_data.csv'),
                                dtype = self.dtyp)
        self.imp = SimpleImputer(missing_values = np.nan, strategy = 'constant', fill_value = 'NA')
//  [python]
def predictor_from_file(config_path, model_path):
        with open(config_path, "rb") as f:
            config_dict = json.load(f)

        if config_dict["model_type"] == "neural_bandit":
            # initialize the pytorch model and put it in `eval` mode
            model = embed_dnn.EmbedDnn(**config_dict["model_spec"])
            model.load_state_dict(torch.load(model_path))
            model.eval()
        else:
            with open(model_path, "rb") as f:
                model = pickle.load(f)

        # initialize transforms
        transforms = {}
        for feature_name, transform_spec in config_dict["transforms"].items():
            if transform_spec is None:
                # id lists don't have transforms
                transform = None
            elif transform_spec["name"] == "StandardScaler":
                transform = preprocessing.StandardScaler()
                transform.mean_ = np.array(transform_spec["mean"])
                transform.scale_ = np.array(transform_spec["scale"])
                transform.var_ = np.array(transform_spec["var"])
            elif transform_spec["name"] == "OneHotEncoder":
                transform = preprocessing.OneHotEncoder()
                transform.sparse = transform_spec["sparse"]
                transform.categories_ = np.array(transform_spec["categories"])
            else:
                raise Exception(
                    f"Don't know how to load transform_spec of type {transform_spec['name']}"
                )
            transforms[feature_name] = transform

        # initialize imputers
        imputers = {}
        for feature_name, imputer_spec in config_dict["imputers"].items():
            if imputer_spec is None:
                # categoricals & id lists don't have imputers
                imputer = None
            else:
                imputer = preprocessing.Imputer()
                imputer.set_params(**imputer_spec["parameters"])
                imputer.statistics_ = np.array(imputer_spec["statistics"])
            imputers[feature_name] = imputer

        return BanditPredictor(
            feature_config=config_dict["feature_config"],
            float_feature_order=config_dict["float_feature_order"],
            id_feature_order=config_dict["id_feature_order"],
            id_feature_str_to_int_map=config_dict["id_feature_str_to_int_map"],
            transforms=transforms,
            imputers=imputers,
            model=model,
            model_type=config_dict["model_type"],
            reward_type=config_dict["reward_type"],
            model_spec=config_dict["model_spec"],
            dense_features_to_use=config_dict["dense_features_to_use"],
        )
//  [python]
def predictor_from_file(config_path, model_path):
        with open(config_path, "rb") as f:
            config_dict = json.load(f)

        if config_dict["model_type"] == "neural_bandit":
            # initialize the pytorch model and put it in `eval` mode
            model = embed_dnn.EmbedDnn(**config_dict["model_spec"])
            model.load_state_dict(torch.load(model_path))
            model.eval()
        else:
            with open(model_path, "rb") as f:
                model = pickle.load(f)

        # initialize transforms
        transforms = {}
        for feature_name, transform_spec in config_dict["transforms"].items():
            if transform_spec is None:
                # id lists don't have transforms
                transform = None
            elif transform_spec["name"] == "StandardScaler":
                transform = preprocessing.StandardScaler()
                transform.mean_ = np.array(transform_spec["mean"])
                transform.scale_ = np.array(transform_spec["scale"])
                transform.var_ = np.array(transform_spec["var"])
            elif transform_spec["name"] == "OneHotEncoder":
                transform = preprocessing.OneHotEncoder()
                transform.sparse = transform_spec["sparse"]
                transform.categories_ = np.array(transform_spec["categories"])
            else:
                raise Exception(
                    f"Don't know how to load transform_spec of type {transform_spec['name']}"
                )
            transforms[feature_name] = transform

        # initialize imputers
        imputers = {}
        for feature_name, imputer_spec in config_dict["imputers"].items():
            if imputer_spec is None:
                # categoricals & id lists don't have imputers
                imputer = None
            else:
                imputer = SimpleImputer()
                imputer.set_params(**imputer_spec["parameters"])
                imputer.statistics_ = np.array(imputer_spec["statistics"])
            imputers[feature_name] = imputer

        return BanditPredictor(
            feature_config=config_dict["feature_config"],
            float_feature_order=config_dict["float_feature_order"],
            id_feature_order=config_dict["id_feature_order"],
            id_feature_str_to_int_map=config_dict["id_feature_str_to_int_map"],
            transforms=transforms,
            imputers=imputers,
            model=model,
            model_type=config_dict["model_type"],
            reward_type=config_dict["reward_type"],
            model_spec=config_dict["model_spec"],
            dense_features_to_use=config_dict["dense_features_to_use"],
        )
//  [python]
def fit(self, data, cat_feats=None, num_feats=None,
          fill_value=-1, n_neighbors=5, **kwargs):

    if cat_feats is None: cat_feats = []
    if num_feats is None: num_feats = []

    assert len(cat_feats + num_feats) != 0, "Please specify \
    categorical and numerical features."

    self._cat_feats = cat_feats
    self._num_feats = num_feats

    df = data.copy()

    ####### REMAINING VARIABLES
    remaining_feats = set(df.columns) - set(cat_feats) - set(num_feats)

    if self.remaining == 'drop':
      df = df.drop(columns=list(remaining_feats))

    ####### CAT VARIABLES
    if self._cat_feats:
      if self.cat_feat_strat == 'replace':
        self._cat_base_imputer = Imputer(strategy='constant',
                                               fill_value=fill_value).fit(df[cat_feats])
      elif self.cat_feat_strat == 'mode':
        self._cat_base_imputer = Imputer(strategy='most_frequent',
                                               fill_value=fill_value).fit(df[cat_feats])

    ####### NUM VARIABLES
    if self._num_feats:
      if self.num_feat_strat == 'mean':
        self._num_base_imputer = Imputer(strategy='mean').fit(df[num_feats])
      elif self.num_feat_strat == 'median':
        self._num_base_imputer = Imputer(strategy='median').fit(df[num_feats])
      elif self.num_feat_strat == 'knn':
        self._num_base_imputer = KNNImputer(n_neighbors=n_neighbors,
                                            **kwargs).fit(df[num_feats])
      elif self.num_feat_strat == 'missforest':
        from missingpy import MissForest
        self._num_base_imputer = MissForest(**kwargs).fit(df[num_feats])

    self.fitted = True
    return self
//  [python]
def fit(self, data, cat_feats=None, num_feats=None,
          fill_value=-1, n_neighbors=5, **kwargs):

    if cat_feats is None: cat_feats = []
    if num_feats is None: num_feats = []

    assert len(cat_feats + num_feats) != 0, "Please specify \
    categorical and numerical features."

    self._cat_feats = cat_feats
    self._num_feats = num_feats

    df = data.copy()

    ####### REMAINING VARIABLES
    remaining_feats = set(df.columns) - set(cat_feats) - set(num_feats)

    if self.remaining == 'drop':
      df = df.drop(columns=list(remaining_feats))

    ####### CAT VARIABLES
    if self._cat_feats:
      if self.cat_feat_strat == 'replace':
        self._cat_base_imputer = SimpleImputer(strategy='constant',
                                               fill_value=fill_value).fit(df[cat_feats])
      elif self.cat_feat_strat == 'mode':
        self._cat_base_imputer = SimpleImputer(strategy='most_frequent',
                                               fill_value=fill_value).fit(df[cat_feats])

    ####### NUM VARIABLES
    if self._num_feats:
      if self.num_feat_strat == 'mean':
        self._num_base_imputer = SimpleImputer(strategy='mean').fit(df[num_feats])
      elif self.num_feat_strat == 'median':
        self._num_base_imputer = SimpleImputer(strategy='median').fit(df[num_feats])
      elif self.num_feat_strat == 'knn':
        self._num_base_imputer = KNNImputer(n_neighbors=n_neighbors,
                                            **kwargs).fit(df[num_feats])
      elif self.num_feat_strat == 'missforest':
        from missingpy import MissForest
        self._num_base_imputer = MissForest(**kwargs).fit(df[num_feats])

    self.fitted = True
    return self
//  [python]
def _get_nested_pipelines_and_data(self, last_transformer=None):
        # returns a pipeline that can be used to test nested pipelines. When last_transformer is not None, it is
        # added as the last transformer in pipeline_1
        steps = [("a", Imputer()), ("b", OneHotEncoder())]
        if last_transformer:
            steps.append(("c", last_transformer))
        pipeline_1 = Pipeline(steps)
        pipeline = Pipeline([("a", Imputer()), ("b", pipeline_1)])
        x = np.zeros((5, 2))
        x[0, 0] = 1
        x[0, 1] = 1
        x[1, 0] = 2

        return pipeline.fit(x), x
//  [python]
def _get_nested_pipelines_and_data(self, last_transformer=None):
        # returns a pipeline that can be used to test nested pipelines. When last_transformer is not None, it is
        # added as the last transformer in pipeline_1
        steps = [("a", SimpleImputer()), ("b", OneHotEncoder())]
        if last_transformer:
            steps.append(("c", last_transformer))
        pipeline_1 = Pipeline(steps)
        pipeline = Pipeline([("a", SimpleImputer()), ("b", pipeline_1)])
        x = np.zeros((5, 2))
        x[0, 0] = 1
        x[0, 1] = 1
        x[1, 0] = 2

        return pipeline.fit(x), x
//  [python]
def test_get_feature_mapper_tuple_for_pipeline(self):
        pipeline = Pipeline([("a", Imputer()), ("b", Imputer()), ("c", OneHotEncoder())])
        x = np.zeros((5, 2))
        x[0, 0] = 1
        x[0, 1] = 1
        pipeline.fit(x)
        feature_mapper = get_feature_mapper_for_pipeline(pipeline)
        feature_mapper.transform(x)

        feature_map = np.zeros((2, 4))
        feature_map[0, :2] = 1
        feature_map[1, 2:] = 1
        assert np.all(feature_mapper.feature_map == feature_map)
//  [python]
def test_get_feature_mapper_tuple_for_pipeline(self):
        pipeline = Pipeline([("a", SimpleImputer()), ("b", SimpleImputer()), ("c", OneHotEncoder())])
        x = np.zeros((5, 2))
        x[0, 0] = 1
        x[0, 1] = 1
        pipeline.fit(x)
        feature_mapper = get_feature_mapper_for_pipeline(pipeline)
        feature_mapper.transform(x)

        feature_map = np.zeros((2, 4))
        feature_map[0, :2] = 1
        feature_map[1, 2:] = 1
        assert np.all(feature_mapper.feature_map == feature_map)
//  [python]
def test_identity_mapper(self):
        x = np.zeros((5, 2))
        imputer = Imputer()
        imputer.fit(x)
        mapper = IdentityMapper(imputer)
        mapper.transform(x)

        feature_map = np.eye(2)
        assert np.all(mapper.feature_map == feature_map)
//  [python]
def test_identity_mapper(self):
        x = np.zeros((5, 2))
        imputer = SimpleImputer()
        imputer.fit(x)
        mapper = IdentityMapper(imputer)
        mapper.transform(x)

        feature_map = np.eye(2)
        assert np.all(mapper.feature_map == feature_map)
//  [python]
def _setup_pipeline(
        self,
        original_features: List[str],
        time_features: List[str],
    ) -> Pipeline:
        """
        Create the pipeline that is later fitted. Includes:
        - BuildRollingFeatures
        - SimpleImputer
        - Model (depending on self.estimator_type)

        Parameters
        ----------
        original_features : list of strings
            list of original feature names
        time_features : list of strings
            list of time feature names

        Returns
        -------
        Pipeline
            resulting pipeline that can be fitted
        """

        rolls = []
        for feature in original_features:
            for i in range(len(self.window_sizes)):
                # only keep originals once to prevent doubling
                rolls.append(
                    (
                        "%s_%d" % (feature, i),
                        BuildRollingFeatures(keep_original=(i == 0)),
                        [feature],
                    )
                )

        for time_feature in time_features:
            rolls.append((time_feature, IdentityFeatureEngineer(), [time_feature]))

        pipeline = Pipeline(
            steps=[
                ("rollpipe", ColumnTransformer(rolls, n_jobs=-1)),
                ("imputer", Imputer()),
                ("model", self._setup_estimator()),
            ]
        )

        return pipeline
//  [python]
def _setup_pipeline(
        self,
        original_features: List[str],
        time_features: List[str],
    ) -> Pipeline:
        """
        Create the pipeline that is later fitted. Includes:
        - BuildRollingFeatures
        - SimpleImputer
        - Model (depending on self.estimator_type)

        Parameters
        ----------
        original_features : list of strings
            list of original feature names
        time_features : list of strings
            list of time feature names

        Returns
        -------
        Pipeline
            resulting pipeline that can be fitted
        """

        rolls = []
        for feature in original_features:
            for i in range(len(self.window_sizes)):
                # only keep originals once to prevent doubling
                rolls.append(
                    (
                        "%s_%d" % (feature, i),
                        BuildRollingFeatures(keep_original=(i == 0)),
                        [feature],
                    )
                )

        for time_feature in time_features:
            rolls.append((time_feature, IdentityFeatureEngineer(), [time_feature]))

        pipeline = Pipeline(
            steps=[
                ("rollpipe", ColumnTransformer(rolls, n_jobs=-1)),
                ("imputer", SimpleImputer()),
                ("model", self._setup_estimator()),
            ]
        )

        return pipeline
//  [python]
def fit(self, X, y=None):
        """Fit OneHotEncoder to X, then transform X.

        Equivalent to self.fit(X).transform(X), but more convenient and more
        efficient. See fit for the parameters, transform for the return value.

        Parameters
        ----------
        X : array-like or sparse matrix, shape=(n_samples, n_features)
            Dense array or sparse matrix.
        y: array-like {n_samples,} (Optional, ignored)
            Feature labels
        """

        if (self.columns == "categorical" or self.columns == "numeric") and not isinstance(X, pd.DataFrame):
            raise ValueError(f"Invalid value for columns: {self.columns}. "
                             "Only 'all' or <list> is supported for np arrays")

        if self.columns == "categorical":
            self.columns_ = list(X.select_dtypes(exclude='number').columns)
        elif self.columns == "numeric":
            self.columns_ =  [col for col in X.columns if is_numeric_dtype(X[col])]
        elif self.columns == "all":
            if isinstance(X, pd.DataFrame):
                self.columns_ = X.columns
            else:
                self.columns_ = list(range(X.shape[1]))
        elif isinstance(self.columns, list):
            self.columns_ = self.columns
        else:
            raise ValueError(f"Invalid value for columns: {self.columns}")
        
        if len(self.columns_) == 0:
            return self
        
        self.imputer = sklearn.preprocessing.Imputer(missing_values=self.missing_values,
                                                    strategy=self.strategy,
                                                    copy=self.copy,
                                                    add_indicator=self.add_indicator)
        
        if isinstance(X, pd.DataFrame):
            self.imputer.set_output(transform="pandas")

        if isinstance(X, pd.DataFrame):
            self.imputer.fit(X[self.columns_], y)
        else:
            self.imputer.fit(X[:, self.columns_], y)

        return self
//  [python]
def fit(self, X, y=None):
        """Fit OneHotEncoder to X, then transform X.

        Equivalent to self.fit(X).transform(X), but more convenient and more
        efficient. See fit for the parameters, transform for the return value.

        Parameters
        ----------
        X : array-like or sparse matrix, shape=(n_samples, n_features)
            Dense array or sparse matrix.
        y: array-like {n_samples,} (Optional, ignored)
            Feature labels
        """

        if (self.columns == "categorical" or self.columns == "numeric") and not isinstance(X, pd.DataFrame):
            raise ValueError(f"Invalid value for columns: {self.columns}. "
                             "Only 'all' or <list> is supported for np arrays")

        if self.columns == "categorical":
            self.columns_ = list(X.select_dtypes(exclude='number').columns)
        elif self.columns == "numeric":
            self.columns_ =  [col for col in X.columns if is_numeric_dtype(X[col])]
        elif self.columns == "all":
            if isinstance(X, pd.DataFrame):
                self.columns_ = X.columns
            else:
                self.columns_ = list(range(X.shape[1]))
        elif isinstance(self.columns, list):
            self.columns_ = self.columns
        else:
            raise ValueError(f"Invalid value for columns: {self.columns}")
        
        if len(self.columns_) == 0:
            return self
        
        self.imputer = sklearn.impute.SimpleImputer(missing_values=self.missing_values,
                                                    strategy=self.strategy,
                                                    fill_value=self.fill_value,
                                                    copy=self.copy,
                                                    add_indicator=self.add_indicator,
                                                    keep_empty_features=self.keep_empty_features)
        
        if isinstance(X, pd.DataFrame):
            self.imputer.set_output(transform="pandas")

        if isinstance(X, pd.DataFrame):
            self.imputer.fit(X[self.columns_], y)
        else:
            self.imputer.fit(X[:, self.columns_], y)

        return self
//  [python]
def define_pipeline(self):
        logger.info("init pipeline")
        self.pipeline = Pipeline(
            [
                (
                    "simple_imputer",
                    Imputer(
                        missing_values=np.nan,
                        strategy="constant",
                        fill_value=0,
                    ),
                ),
                (
                    "min_max_scaler",
                    MinMaxScaler(),
                ),
            ]
        )

        logger.info(f"pipeline: {self.pipeline}")
//  [python]
def define_pipeline(self):
        logger.info("init pipeline")
        self.pipeline = Pipeline(
            [
                (
                    "simple_imputer",
                    SimpleImputer(
                        missing_values=np.nan,
                        strategy="constant",
                        fill_value=0,
                    ),
                ),
                (
                    "min_max_scaler",
                    MinMaxScaler(),
                ),
            ]
        )

        logger.info(f"pipeline: {self.pipeline}")
//  [python]
X, y = fetch_openml(name="miceprotein", return_X_y=True)
# Fill missing values with the mean
X = Imputer().fit_transform(X)
//  [python]
X, y = fetch_openml(name="miceprotein", return_X_y=True)
# Fill missing values with the mean
X = SimpleImputer().fit_transform(X)
//  [python]
idx = 30000
X_train, Y_train = data[:idx, :], labels[:idx]
X_test, Y_test = data[idx + 1:, :], labels[idx + 1:]

# feature transformation pipeline
ordinal_features = [x for x in range(len(feature_names)) if x not in list(category_map.keys())]
ordinal_transformer = Pipeline(steps=[('imputer', Imputer(strategy='median')),
                                      ('scaler', StandardScaler())])
//  [python]
idx = 30000
X_train, Y_train = data[:idx, :], labels[:idx]
X_test, Y_test = data[idx + 1:, :], labels[idx + 1:]

# feature transformation pipeline
ordinal_features = [x for x in range(len(feature_names)) if x not in list(category_map.keys())]
ordinal_transformer = Pipeline(steps=[('imputer', SimpleImputer(strategy='median')),
                                      ('scaler', StandardScaler())])
//  [python]
def impute(self, series, strategy, fill_value):
        """
        Impute missing values in a series
        """
        from sklearn.preprocessing import Imputer
        imputer = Imputer(strategy=strategy, fill_value=fill_value)
        series_fit = series.dropna()
        if str(series.dtype) in self.constants.OBJECT_TYPES:
            series_fit = series_fit.astype(str)
        values = series_fit.values.reshape(-1, 1)
        if len(values):
            imputer.fit(values)
            return imputer.transform(series.fillna(np.nan).values.reshape(-1, 1))
        else:
            logger.warn("list to fit imputer is empty, try cols.fill_na instead.")
            return series
//  [python]
def impute(self, series, strategy, fill_value):
        """
        Impute missing values in a series
        """
        from sklearn.impute import SimpleImputer
        imputer = SimpleImputer(strategy=strategy, fill_value=fill_value)
        series_fit = series.dropna()
        if str(series.dtype) in self.constants.OBJECT_TYPES:
            series_fit = series_fit.astype(str)
        values = series_fit.values.reshape(-1, 1)
        if len(values):
            imputer.fit(values)
            return imputer.transform(series.fillna(np.nan).values.reshape(-1, 1))
        else:
            logger.warn("list to fit imputer is empty, try cols.fill_na instead.")
            return series
//  [python]
def generate_rule_dict(missing_value, data, strategy):
    # treat all missing_value as missing even if they were not
    data = data.replace({missing_value: np.nan})
    # this step will unify all np.nan, pd.NA and None values
    # just like these NA or null values entered arrow and all becoming null
    data = data.fillna(np.nan)

    imputer = Imputer(missing_values=np.nan, strategy=strategy)
    imputer.fit(data)
    fill_value_dict = {
        k: v for k, v in zip(imputer.feature_names_in_, imputer.statistics_)
    }
    fillna_rules = {"missing_value": missing_value, "fill_value_rule": fill_value_dict}
    return fillna_rules
//  [python]
def generate_rule_dict(missing_value, data, strategy):
    # treat all missing_value as missing even if they were not
    data = data.replace({missing_value: np.nan})
    # this step will unify all np.nan, pd.NA and None values
    # just like these NA or null values entered arrow and all becoming null
    data = data.fillna(np.nan)

    imputer = SimpleImputer(missing_values=np.nan, strategy=strategy)
    imputer.fit(data)
    fill_value_dict = {
        k: v for k, v in zip(imputer.feature_names_in_, imputer.statistics_)
    }
    fillna_rules = {"missing_value": missing_value, "fill_value_rule": fill_value_dict}
    return fillna_rules
//  [python]
def make_logreg(categorical_features):
    cont_pipe = make_pipeline(StandardScaler(), Imputer())
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore'), categorical_features), remainder=cont_pipe)
    return make_pipeline(preprocess, LogisticRegression(max_iter=1000))
//  [python]
def make_logreg(categorical_features):
    cont_pipe = make_pipeline(StandardScaler(), SimpleImputer())
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore'), categorical_features), remainder=cont_pipe)
    return make_pipeline(preprocess, LogisticRegression(max_iter=1000))
//  [python]
def make_knn(categorical_features):
    cont_pipe = make_pipeline(StandardScaler(), Imputer())
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore', sparse_output=False), categorical_features), remainder=cont_pipe)
    return make_pipeline(preprocess, KNeighborsClassifier())
//  [python]
def make_knn(categorical_features):
    cont_pipe = make_pipeline(StandardScaler(), SimpleImputer())
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore', sparse_output=False), categorical_features), remainder=cont_pipe)
    return make_pipeline(preprocess, KNeighborsClassifier())
//  [python]
def make_rf(categorical_features):
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore', sparse_output=False), categorical_features), remainder=Imputer())
    return make_pipeline(preprocess, RandomForestClassifier())
//  [python]
def make_rf(categorical_features):
    preprocess = make_column_transformer((OneHotEncoder(handle_unknown='ignore', sparse_output=False), categorical_features), remainder=SimpleImputer())
    return make_pipeline(preprocess, RandomForestClassifier())
//  [python]
def get_inference_pipeline(args):
    # We need 3 separate preprocessing "tracks":
    # - one for categorical features
    # - one for numerical features
    # - one for textual ("nlp") features
    # Categorical preprocessing pipeline
    categorical_features = sorted(["time_signature", "key"])
    categorical_transformer = make_pipeline(
        Imputer(strategy="constant", fill_value=0), OrdinalEncoder()
    )
    # Numerical preprocessing pipeline
    numeric_features = sorted([
        "danceability",
        "energy",
        "loudness",
        "speechiness",
        "acousticness",
        "instrumentalness",
        "liveness",
        "valence",
        "tempo",
        "duration_ms",
    ])
    numeric_transformer = make_pipeline(
        Imputer(strategy="median"), StandardScaler()
    )
    # Textual ("nlp") preprocessing pipeline
    nlp_features = ["text_feature"]
    # This trick is needed because Imputer wants a 2d input, but
    # TfidfVectorizer wants a 1d input. So we reshape in between the two steps
    reshape_to_1d = FunctionTransformer(np.reshape, kw_args={"newshape": -1})
    nlp_transformer = make_pipeline(
        Imputer(strategy="constant", fill_value=""),
        reshape_to_1d,
        TfidfVectorizer(binary=True),
    )
    # Put the 3 tracks together into one pipeline using the ColumnTransformer
    # This also drops the columns that we are not explicitly transforming
    preprocessor = ColumnTransformer(
        transformers=[
            ("num", numeric_transformer, numeric_features),
            ("cat", categorical_transformer, categorical_features),
            ("nlp1", nlp_transformer, nlp_features),
        ],
        remainder="drop",  # This drops the columns that we do not transform
    )
    # Get the configuration for the model
    with open(args.model_config) as fp:
        model_config = json.load(fp)
    # Add it to the W&B configuration so the values for the hyperparams
    # are tracked
    wandb.config.update(model_config)
    # Append classifier to preprocessing pipeline.
    # Now we have a full prediction pipeline.
    pipe = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("classifier", RandomForestClassifier(**model_config)),
        ]
    )
    return pipe
//  [python]
def get_inference_pipeline(args):
    # We need 3 separate preprocessing "tracks":
    # - one for categorical features
    # - one for numerical features
    # - one for textual ("nlp") features
    # Categorical preprocessing pipeline
    categorical_features = sorted(["time_signature", "key"])
    categorical_transformer = make_pipeline(
        SimpleImputer(strategy="constant", fill_value=0), OrdinalEncoder()
    )
    # Numerical preprocessing pipeline
    numeric_features = sorted([
        "danceability",
        "energy",
        "loudness",
        "speechiness",
        "acousticness",
        "instrumentalness",
        "liveness",
        "valence",
        "tempo",
        "duration_ms",
    ])
    numeric_transformer = make_pipeline(
        SimpleImputer(strategy="median"), StandardScaler()
    )
    # Textual ("nlp") preprocessing pipeline
    nlp_features = ["text_feature"]
    # This trick is needed because SimpleImputer wants a 2d input, but
    # TfidfVectorizer wants a 1d input. So we reshape in between the two steps
    reshape_to_1d = FunctionTransformer(np.reshape, kw_args={"newshape": -1})
    nlp_transformer = make_pipeline(
        SimpleImputer(strategy="constant", fill_value=""),
        reshape_to_1d,
        TfidfVectorizer(binary=True),
    )
    # Put the 3 tracks together into one pipeline using the ColumnTransformer
    # This also drops the columns that we are not explicitly transforming
    preprocessor = ColumnTransformer(
        transformers=[
            ("num", numeric_transformer, numeric_features),
            ("cat", categorical_transformer, categorical_features),
            ("nlp1", nlp_transformer, nlp_features),
        ],
        remainder="drop",  # This drops the columns that we do not transform
    )
    # Get the configuration for the model
    with open(args.model_config) as fp:
        model_config = json.load(fp)
    # Add it to the W&B configuration so the values for the hyperparams
    # are tracked
    wandb.config.update(model_config)
    # Append classifier to preprocessing pipeline.
    # Now we have a full prediction pipeline.
    pipe = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("classifier", RandomForestClassifier(**model_config)),
        ]
    )
    return pipe
//  [python]
def test_pipeline_compatibility() -> None:
    """Check that MAPIE works on pipeline based on pandas dataframes"""
    X = pd.DataFrame(
        {
            "x_cat": ["A", "A", "B", "A", "A", "B"],
            "x_num": [0, 1, 1, 4, np.nan, 5],
        }
    )
    y = pd.Series([0, 1, 2, 0, 1, 0])
    numeric_preprocessor = Pipeline(
        [
            ("imputer", Imputer(strategy="mean")),
        ]
    )
    categorical_preprocessor = Pipeline(
        steps=[
            ("encoding", OneHotEncoder(handle_unknown="ignore"))
        ]
    )
    preprocessor = ColumnTransformer(
        [
            ("cat", categorical_preprocessor, ["x_cat"]),
            ("num", numeric_preprocessor, ["x_num"])
        ]
    )
    pipe = make_pipeline(preprocessor, LogisticRegression())
    pipe.fit(X, y)
    mapie = MapieCalibrator(estimator=pipe)
    mapie.fit(X, y)
    mapie.predict(X)
//  [python]
def test_pipeline_compatibility() -> None:
    """Check that MAPIE works on pipeline based on pandas dataframes"""
    X = pd.DataFrame(
        {
            "x_cat": ["A", "A", "B", "A", "A", "B"],
            "x_num": [0, 1, 1, 4, np.nan, 5],
        }
    )
    y = pd.Series([0, 1, 2, 0, 1, 0])
    numeric_preprocessor = Pipeline(
        [
            ("imputer", SimpleImputer(strategy="mean")),
        ]
    )
    categorical_preprocessor = Pipeline(
        steps=[
            ("encoding", OneHotEncoder(handle_unknown="ignore"))
        ]
    )
    preprocessor = ColumnTransformer(
        [
            ("cat", categorical_preprocessor, ["x_cat"]),
            ("num", numeric_preprocessor, ["x_num"])
        ]
    )
    pipe = make_pipeline(preprocessor, LogisticRegression())
    pipe.fit(X, y)
    mapie = MapieCalibrator(estimator=pipe)
    mapie.fit(X, y)
    mapie.predict(X)