jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _load_support_dataset():
  """Helper function to load and preprocess the SUPPORT dataset.

  The SUPPORT Dataset comes from the Vanderbilt University study
  to estimate survival for seriously ill hospitalized adults [1].

  Please refer to http://biostat.mc.vanderbilt.edu/wiki/Main/SupportDesc.
  for the original datasource.

  References
  ----------
  [1]: Knaus WA, Harrell FE, Lynn J et al. (1995): The SUPPORT prognostic
  model: Objective estimates of survival for seriously ill hospitalized
  adults. Annals of Internal Medicine 122:191-203.

  """

  data = pkgutil.get_data(__name__, 'datasets/support2.csv')
  data = pd.read_csv(io.BytesIO(data))
  x1 = data[['age', 'num.co', 'meanbp', 'wblc', 'hrt', 'resp', 'temp',
             'pafi', 'alb', 'bili', 'crea', 'sod', 'ph', 'glucose', 'bun',
             'urine', 'adlp', 'adls']]

  catfeats = ['sex', 'dzgroup', 'dzclass', 'income', 'race', 'ca']
  x2 = pd.get_dummies(data[catfeats])

  x = np.concatenate([x1, x2], axis=1)
  t = data['d.time'].values
  e = data['death'].values

  x = SimpleImputer(missing_values=np.nan, strategy='mean').fit_transform(x)
  x = StandardScaler().fit_transform(x)

  remove = ~np.isnan(t)
  return x[remove], t[remove], e[remove]
//  [python]
def _load_framingham_dataset(sequential):
  """Helper function to load and preprocess the Framingham dataset.

  The Framingham Dataset is a subset of 4,434 participants of the well known,
  ongoing Framingham Heart study [1] for studying epidemiology for
  hypertensive and arteriosclerotic cardiovascular disease. It is a popular
  dataset for longitudinal survival analysis with time dependent covariates.

  Parameters
  ----------
  sequential: bool
    If True returns a list of np.arrays for each individual.
    else, returns collapsed results for each time step. To train
    recurrent neural models you would typically use True.

  References
  ----------
  [1] Dawber, Thomas R., Gilcin F. Meadors, and Felix E. Moore Jr.
  "Epidemiological approaches to heart disease: the Framingham Study."
  American Journal of Public Health and the Nations Health 41.3 (1951).

  """

  data = pkgutil.get_data(__name__, 'datasets/framingham.csv')
  data = pd.read_csv(io.BytesIO(data))

  dat_cat = data[['SEX', 'CURSMOKE', 'DIABETES', 'BPMEDS',
                  'educ', 'PREVCHD', 'PREVAP', 'PREVMI',
                  'PREVSTRK', 'PREVHYP']]
  dat_num = data[['TOTCHOL', 'AGE', 'SYSBP', 'DIABP',
                  'CIGPDAY', 'BMI', 'HEARTRTE', 'GLUCOSE']]

  x1 = pd.get_dummies(dat_cat).values
  x2 = dat_num.values
  x = np.hstack([x1, x2])

  time = (data['TIMEDTH'] - data['TIME']).values
  event = data['DEATH'].values

  x = Imputer(missing_values=np.nan, strategy='mean').fit_transform(x)
  x_ = StandardScaler().fit_transform(x)

  if not sequential:
    return x_, time, event
  else:
    x, t, e = [], [], []
    for id_ in sorted(list(set(data['RANDID']))):
      x.append(x_[data['RANDID'] == id_])
      t.append(time[data['RANDID'] == id_])
      e.append(event[data['RANDID'] == id_])
    return x, t, e
//  [python]
def _load_framingham_dataset(sequential):
  """Helper function to load and preprocess the Framingham dataset.

  The Framingham Dataset is a subset of 4,434 participants of the well known,
  ongoing Framingham Heart study [1] for studying epidemiology for
  hypertensive and arteriosclerotic cardiovascular disease. It is a popular
  dataset for longitudinal survival analysis with time dependent covariates.

  Parameters
  ----------
  sequential: bool
    If True returns a list of np.arrays for each individual.
    else, returns collapsed results for each time step. To train
    recurrent neural models you would typically use True.

  References
  ----------
  [1] Dawber, Thomas R., Gilcin F. Meadors, and Felix E. Moore Jr.
  "Epidemiological approaches to heart disease: the Framingham Study."
  American Journal of Public Health and the Nations Health 41.3 (1951).

  """

  data = pkgutil.get_data(__name__, 'datasets/framingham.csv')
  data = pd.read_csv(io.BytesIO(data))

  dat_cat = data[['SEX', 'CURSMOKE', 'DIABETES', 'BPMEDS',
                  'educ', 'PREVCHD', 'PREVAP', 'PREVMI',
                  'PREVSTRK', 'PREVHYP']]
  dat_num = data[['TOTCHOL', 'AGE', 'SYSBP', 'DIABP',
                  'CIGPDAY', 'BMI', 'HEARTRTE', 'GLUCOSE']]

  x1 = pd.get_dummies(dat_cat).values
  x2 = dat_num.values
  x = np.hstack([x1, x2])

  time = (data['TIMEDTH'] - data['TIME']).values
  event = data['DEATH'].values

  x = SimpleImputer(missing_values=np.nan, strategy='mean').fit_transform(x)
  x_ = StandardScaler().fit_transform(x)

  if not sequential:
    return x_, time, event
  else:
    x, t, e = [], [], []
    for id_ in sorted(list(set(data['RANDID']))):
      x.append(x_[data['RANDID'] == id_])
      t.append(time[data['RANDID'] == id_])
      e.append(event[data['RANDID'] == id_])
    return x, t, e
//  [python]
def build_prep(num_columns, cat_columns):
    # For categorical features: impute missing values as 'UNK', onehot encode
    cat_pipe = make_pipeline(
        Imputer(strategy="constant", fill_value="UNK", verbose=1),
        OneHotEncoder(sparse=False, handle_unknown="ignore"),  # May lead to silent failure!
    )
    # For numerical features: impute missing values as -1
    num_pipe = make_pipeline(
        Imputer(strategy="constant", fill_value=-1, verbose=1),
    )

    coltf = ColumnTransformer(transformers=[("num", num_pipe, num_columns), ("cat", cat_pipe, cat_columns)])
    return coltf
//  [python]
def build_prep(num_columns, cat_columns):
    # For categorical features: impute missing values as 'UNK', onehot encode
    cat_pipe = make_pipeline(
        SimpleImputer(strategy="constant", fill_value="UNK", verbose=1),
        OneHotEncoder(sparse=False, handle_unknown="ignore"),  # May lead to silent failure!
    )
    # For numerical features: impute missing values as -1
    num_pipe = make_pipeline(
        SimpleImputer(strategy="constant", fill_value=-1, verbose=1),
    )

    coltf = ColumnTransformer(transformers=[("num", num_pipe, num_columns), ("cat", cat_pipe, cat_columns)])
    return coltf
//  [python]
def _load_pbc_dataset(sequential):
  """Helper function to load and preprocess the PBC dataset

  The Primary biliary cirrhosis (PBC) Dataset [1] is well known
  dataset for evaluating survival analysis models with time
  dependent covariates.

  Parameters
  ----------
  sequential: bool
    If True returns a list of np.arrays for each individual.
    else, returns collapsed results for each time step. To train
    recurrent neural models you would typically use True.

  References
  ----------
  [1] Fleming, Thomas R., and David P. Harrington. Counting processes and
  survival analysis. Vol. 169. John Wiley & Sons, 2011.

  """

  data = pkgutil.get_data(__name__, 'datasets/pbc2.csv')
  data = pd.read_csv(io.BytesIO(data))

  data['histologic'] = data['histologic'].astype(str)
  dat_cat = data[['drug', 'sex', 'ascites', 'hepatomegaly',
                  'spiders', 'edema', 'histologic']]
  dat_num = data[['serBilir', 'serChol', 'albumin', 'alkaline',
                  'SGOT', 'platelets', 'prothrombin']]
  age = data['age'] + data['years']

  x1 = pd.get_dummies(dat_cat).values
  x2 = dat_num.values
  x3 = age.values.reshape(-1, 1)
  x = np.hstack([x1, x2, x3])

  time = (data['years'] - data['year']).values
  event = data['status2'].values

  x = Imputer(missing_values=np.nan, strategy='mean').fit_transform(x)
  x_ = StandardScaler().fit_transform(x)

  if not sequential:
    return x_, time, event
  else:
    x, t, e = [], [], []
    for id_ in sorted(list(set(data['id']))):
      x.append(x_[data['id'] == id_])
      t.append(time[data['id'] == id_])
      e.append(event[data['id'] == id_])
    return x, t, e
//  [python]
def _load_pbc_dataset(sequential):
  """Helper function to load and preprocess the PBC dataset

  The Primary biliary cirrhosis (PBC) Dataset [1] is well known
  dataset for evaluating survival analysis models with time
  dependent covariates.

  Parameters
  ----------
  sequential: bool
    If True returns a list of np.arrays for each individual.
    else, returns collapsed results for each time step. To train
    recurrent neural models you would typically use True.

  References
  ----------
  [1] Fleming, Thomas R., and David P. Harrington. Counting processes and
  survival analysis. Vol. 169. John Wiley & Sons, 2011.

  """

  data = pkgutil.get_data(__name__, 'datasets/pbc2.csv')
  data = pd.read_csv(io.BytesIO(data))

  data['histologic'] = data['histologic'].astype(str)
  dat_cat = data[['drug', 'sex', 'ascites', 'hepatomegaly',
                  'spiders', 'edema', 'histologic']]
  dat_num = data[['serBilir', 'serChol', 'albumin', 'alkaline',
                  'SGOT', 'platelets', 'prothrombin']]
  age = data['age'] + data['years']

  x1 = pd.get_dummies(dat_cat).values
  x2 = dat_num.values
  x3 = age.values.reshape(-1, 1)
  x = np.hstack([x1, x2, x3])

  time = (data['years'] - data['year']).values
  event = data['status2'].values

  x = SimpleImputer(missing_values=np.nan, strategy='mean').fit_transform(x)
  x_ = StandardScaler().fit_transform(x)

  if not sequential:
    return x_, time, event
  else:
    x, t, e = [], [], []
    for id_ in sorted(list(set(data['id']))):
      x.append(x_[data['id'] == id_])
      t.append(time[data['id'] == id_])
      e.append(event[data['id'] == id_])
    return x, t, e
//  [python]
data = fetch_adult()
X = data.data
y = (data.target == ">50K").astype(int)
X_train, X_test, y_train_true, y_test_true = train_test_split(
    X, y, test_size=0.33, random_state=42
)

numeric_transformer = Pipeline(
    steps=[
        ("impute", Imputer()),
        ("scaler", StandardScaler()),
    ]
)
//  [python]
data = fetch_adult()
X = data.data
y = (data.target == ">50K").astype(int)
X_train, X_test, y_train_true, y_test_true = train_test_split(
    X, y, test_size=0.33, random_state=42
)

numeric_transformer = Pipeline(
    steps=[
        ("impute", SimpleImputer()),
        ("scaler", StandardScaler()),
    ]
)
//  [python]
def get_models(horizon):
    model_class = PlayingChanceModel
    return {
        **{
            f"lasso_{alpha}": model_class(
                model=make_pipeline(
                    Imputer(), StandardScaler(), LogisticRegression(C=alpha)
                ),
                horizon=horizon,
            )
            for alpha in [0.01, 0.1, 1, 10, 100]
        },
        **{
            f"lasso_poly_{alpha}": model_class(
                model=make_pipeline(
                    Imputer(),
                    PolynomialFeatures(),
                    StandardScaler(),
                    LogisticRegression(C=alpha),
                ),
                horizon=horizon,
            )
            for alpha in [0.01, 0.1, 1, 10]
        },
    }
//  [python]
def get_models(horizon):
    model_class = PlayingChanceModel
    return {
        **{
            f"lasso_{alpha}": model_class(
                model=make_pipeline(
                    SimpleImputer(), StandardScaler(), LogisticRegression(C=alpha)
                ),
                horizon=horizon,
            )
            for alpha in [0.01, 0.1, 1, 10, 100]
        },
        **{
            f"lasso_poly_{alpha}": model_class(
                model=make_pipeline(
                    SimpleImputer(),
                    PolynomialFeatures(),
                    StandardScaler(),
                    LogisticRegression(C=alpha),
                ),
                horizon=horizon,
            )
            for alpha in [0.01, 0.1, 1, 10]
        },
    }
//  [python]
def _fit_element(self, df: pd.DataFrame, col: str = "__all__", ngroup: int = 0) -> Any:
        """
        Fits the imputer on `df`, at the group and/or column level depending onself.groups and
        self.columnwise.

        Parameters
        ----------
        df : pd.DataFrame
            Dataframe on which the imputer is fitted
        col : str, optional
            Column on which the imputer is fitted, by default "__all__"
        ngroup : int, optional
            Id of the group on which the method is applied

        Returns
        -------
        Any
            Return fitted KNN model

        Raises
        ------
        NotDataFrame
            Input has to be a pandas.DataFrame.
        """
        if pd.api.types.is_numeric_dtype(df[col]):
            model = s.Imputer(strategy=self.strategy)
        else:
            model = s.Imputer(strategy="most_frequent")
        return model.fit(df[[col]])
//  [python]
def _fit_element(self, df: pd.DataFrame, col: str = "__all__", ngroup: int = 0) -> Any:
        """
        Fits the imputer on `df`, at the group and/or column level depending onself.groups and
        self.columnwise.

        Parameters
        ----------
        df : pd.DataFrame
            Dataframe on which the imputer is fitted
        col : str, optional
            Column on which the imputer is fitted, by default "__all__"
        ngroup : int, optional
            Id of the group on which the method is applied

        Returns
        -------
        Any
            Return fitted KNN model

        Raises
        ------
        NotDataFrame
            Input has to be a pandas.DataFrame.
        """
        if pd.api.types.is_numeric_dtype(df[col]):
            model = skl.impute.SimpleImputer(strategy=self.strategy)
        else:
            model = skl.impute.SimpleImputer(strategy="most_frequent")
        return model.fit(df[[col]])
//  [python]
def preprocess(dataframe, specification):

    X = specification['problem']['predictors']
    y = specification['problem']['targets'][0]

    categorical_features = [i for i in specification['problem']['categorical'] if i != y and i in X]

    categorical_transformer = Pipeline(steps=[
        ('imputer', Imputer(strategy='constant', fill_value='missing')),
        ('onehot', OneHotEncoder(handle_unknown='ignore'))
    ])

    numerical_features = [i for i in X if i not in categorical_features]
    numerical_transformer = Pipeline(steps=[
        ('imputer', Imputer(strategy='median')),
        ('scaler', StandardScaler())
    ])

    preprocessor = ColumnTransformer(transformers=[
        ('numeric', numerical_transformer, numerical_features),
        ('categorical', categorical_transformer, categorical_features)
    ])

    stimulus = dataframe[X]
    stimulus = preprocessor.fit_transform(stimulus)

    return stimulus, preprocessor
//  [python]
def preprocess(dataframe, specification):

    X = specification['problem']['predictors']
    y = specification['problem']['targets'][0]

    categorical_features = [i for i in specification['problem']['categorical'] if i != y and i in X]

    categorical_transformer = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='constant', fill_value='missing')),
        ('onehot', OneHotEncoder(handle_unknown='ignore'))
    ])

    numerical_features = [i for i in X if i not in categorical_features]
    numerical_transformer = Pipeline(steps=[
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', StandardScaler())
    ])

    preprocessor = ColumnTransformer(transformers=[
        ('numeric', numerical_transformer, numerical_features),
        ('categorical', categorical_transformer, categorical_features)
    ])

    stimulus = dataframe[X]
    stimulus = preprocessor.fit_transform(stimulus)

    return stimulus, preprocessor
//  [python]
def test_pipeline_cv():
    X = pd.DataFrame(
        {
            "A": ["a", "b", "c", np.nan, "e", "f"],
            "B": ["z", "y", "x", "w", "v", "u"],
            "C": np.random.normal(100, 20, 6),
        }
    )
    y = np.array([1, 0, 1, 0, 1, 0])
    numeric_vars = ["C"]
    categorical_vars = ["A", "B"]
    ee = EmbeddingEncoder(task="classification", epochs=1)
    cat_transformer = make_pipeline(Imputer(strategy="most_frequent"), ee)
    preprocessor = ColumnTransformer(
        [
            ("num", StandardScaler(), numeric_vars),
            ("cat", cat_transformer, categorical_vars),
        ]
    )
    pipeline = make_pipeline(preprocessor, LogisticRegression())
    param_grid = {
        "columntransformer__cat__embeddingencoder__layers_units": [
            [64, 32, 16],
            [16, 8],
        ]
    }
    cv = GridSearchCV(pipeline, param_grid, cv=2)
    cv.fit(X, y)
    assert hasattr(cv, "best_estimator_")
//  [python]
def test_pipeline_cv():
    X = pd.DataFrame(
        {
            "A": ["a", "b", "c", np.nan, "e", "f"],
            "B": ["z", "y", "x", "w", "v", "u"],
            "C": np.random.normal(100, 20, 6),
        }
    )
    y = np.array([1, 0, 1, 0, 1, 0])
    numeric_vars = ["C"]
    categorical_vars = ["A", "B"]
    ee = EmbeddingEncoder(task="classification", epochs=1)
    cat_transformer = make_pipeline(SimpleImputer(strategy="most_frequent"), ee)
    preprocessor = ColumnTransformer(
        [
            ("num", StandardScaler(), numeric_vars),
            ("cat", cat_transformer, categorical_vars),
        ]
    )
    pipeline = make_pipeline(preprocessor, LogisticRegression())
    param_grid = {
        "columntransformer__cat__embeddingencoder__layers_units": [
            [64, 32, 16],
            [16, 8],
        ]
    }
    cv = GridSearchCV(pipeline, param_grid, cv=2)
    cv.fit(X, y)
    assert hasattr(cv, "best_estimator_")
//  [python]
def fit(self, X: pd.DataFrame):
        """Fit scaler based on given dataframe."""
        self.not_cat_columns = list(set(X.columns) - set(self.cat_columns))

        # SimpleImputer doesn't work on all nan-columns, so first replace them  # noqa: SC100
        X = X.apply(ScaledNumerics._impute_whole_series_to_zero, axis=0)

        if self.cat_columns:
            self.categorical_imputer = Imputer(strategy='most_frequent')
            self.categorical_imputer.fit(X[self.cat_columns])
        if self.not_cat_columns:
            self.numeric_imputer = Imputer(strategy='mean')
            self.numeric_imputer.fit(X[self.not_cat_columns])
            self.scaler = MinMaxScaler()
            self.scaler.fit(X[self.not_cat_columns])

        # Replace non-common categories with special value:
        self.rare_category_encoder = RareCategoryEncoder(max_num_categories=self.max_num_categories,
                                                         cols=self.cat_columns)
        self.rare_category_encoder.fit(X)

        # One-hot encode categorical features:
        self.one_hot_encoder = OneHotEncoder(cols=self.cat_columns, use_cat_names=True)
        self.one_hot_encoder.fit(X)
//  [python]
def fit(self, X: pd.DataFrame):
        """Fit scaler based on given dataframe."""
        self.not_cat_columns = list(set(X.columns) - set(self.cat_columns))

        # SimpleImputer doesn't work on all nan-columns, so first replace them  # noqa: SC100
        X = X.apply(ScaledNumerics._impute_whole_series_to_zero, axis=0)

        if self.cat_columns:
            self.categorical_imputer = SimpleImputer(strategy='most_frequent')
            self.categorical_imputer.fit(X[self.cat_columns])
        if self.not_cat_columns:
            self.numeric_imputer = SimpleImputer(strategy='mean')
            self.numeric_imputer.fit(X[self.not_cat_columns])
            self.scaler = MinMaxScaler()
            self.scaler.fit(X[self.not_cat_columns])

        # Replace non-common categories with special value:
        self.rare_category_encoder = RareCategoryEncoder(max_num_categories=self.max_num_categories,
                                                         cols=self.cat_columns)
        self.rare_category_encoder.fit(X)

        # One-hot encode categorical features:
        self.one_hot_encoder = OneHotEncoder(cols=self.cat_columns, use_cat_names=True)
        self.one_hot_encoder.fit(X)
//  [python]
def disparity2depth_psm(disparity):
    baseline = 0.202993
    focal = 2355.722801
    depth = 250*np.ones(disparity.shape)
    # In SGM there are NAN values, as areas could not be calculated. The SimpleImputer interpolates those wholes.
    # Larger wholes in some cases can not be closed. Check this publication for better results: In Defense of Classical Image Processing: Fast Depth Completion on the CPU
    imp = Imputer(missing_values=np.nan, strategy="mean")
    disparity = imp.fit_transform(disparity)
    depth[disparity == 0] = 250
    depth[disparity != 0] = focal * baseline / disparity[disparity != 0]

    return np.clip(depth,0,250)
//  [python]
def disparity2depth_psm(disparity):
    baseline = 0.202993
    focal = 2355.722801
    depth = 250*np.ones(disparity.shape)
    # In SGM there are NAN values, as areas could not be calculated. The SimpleImputer interpolates those wholes.
    # Larger wholes in some cases can not be closed. Check this publication for better results: In Defense of Classical Image Processing: Fast Depth Completion on the CPU
    imp = SimpleImputer(missing_values=np.nan, strategy="mean")
    disparity = imp.fit_transform(disparity)
    depth[disparity == 0] = 250
    depth[disparity != 0] = focal * baseline / disparity[disparity != 0]

    return np.clip(depth,0,250)
//  [python]
def transform(self, X):
        imputer = Imputer(strategy="mean")
        X['Age'] = imputer.fit_transform(X[['Age']])
        return X
//  [python]
def transform(self, X):
        imputer = SimpleImputer(strategy="mean")
        X['Age'] = imputer.fit_transform(X[['Age']])
        return X
//  [python]
def imputer(data: pd.DataFrame, method: str) -> tuple[dict, np.ndarray]:
    """Apply imputation on missing values.

    Parameters
    ----------
    data : pd.DataFrame
        The dataset with missing values.

    method : str
        The imputation method.

    Returns
    -------
    imputation_config : dict
        The imputation configuration.

    data_imputed : np.ndarray
        The dataset after imputing.
    """
    method2option = {
        "Mean Value": "mean",
        "Median Value": "median",
        "Most Frequent Value": "most_frequent",
        "Constant(Specified Value)": "constant",
    }[method]
    if method2option == "constant":
        filled_value = float_input(0, SECTION[2], "@Specified Value: ")
        imputer = Imputer(missing_values=np.nan, strategy=method2option, fill_value=filled_value)
    else:
        imputer = Imputer(missing_values=np.nan, strategy=method2option)
    print(f"Successfully fill the missing values with the {method2option} value " f"of each feature column respectively.")
    data_imputed = imputer.fit_transform(data)
    imputation_config = {type(imputer).__name__: imputer.get_params()}
    return imputation_config, data_imputed
//  [python]
def imputer(data: pd.DataFrame, method: str) -> tuple[dict, np.ndarray]:
    """Apply imputation on missing values.

    Parameters
    ----------
    data : pd.DataFrame
        The dataset with missing values.

    method : str
        The imputation method.

    Returns
    -------
    imputation_config : dict
        The imputation configuration.

    data_imputed : np.ndarray
        The dataset after imputing.
    """
    method2option = {
        "Mean Value": "mean",
        "Median Value": "median",
        "Most Frequent Value": "most_frequent",
        "Constant(Specified Value)": "constant",
    }[method]
    if method2option == "constant":
        filled_value = float_input(0, SECTION[2], "@Specified Value: ")
        imputer = SimpleImputer(missing_values=np.nan, strategy=method2option, fill_value=filled_value)
    else:
        imputer = SimpleImputer(missing_values=np.nan, strategy=method2option)
    print(f"Successfully fill the missing values with the {method2option} value " f"of each feature column respectively.")
    data_imputed = imputer.fit_transform(data)
    imputation_config = {type(imputer).__name__: imputer.get_params()}
    return imputation_config, data_imputed
//  [python]
def get_training_inference_pipeline(args):

    # Get the configuration for the pipeline
    with open(args.model_config) as fp:
        model_config = yaml.safe_load(fp)
    # Add it to the W&B configuration so the values for the hyperparams
    # are tracked
    wandb.config.update(model_config)

    # 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(model_config["features"]["categorical"])
    categorical_transformer = make_pipeline(
        Imputer(strategy="constant", fill_value=0), OrdinalEncoder()
    )
    # Numerical preprocessing pipeline
    numeric_features = sorted(model_config["features"]["numerical"])
    numeric_transformer = make_pipeline(
        Imputer(strategy="median"), StandardScaler()
    )
    # Textual ("nlp") preprocessing pipeline
    nlp_features = sorted(model_config["features"]["nlp"])
    # 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, max_features=model_config["tfidf"]["max_features"]
        ),
    )
    # 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
    )

    # Append classifier to preprocessing pipeline.
    # Now we have a full prediction pipeline.
    pipe = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("classifier", RandomForestClassifier(**model_config["random_forest"])),
        ]
    )
    return pipe
//  [python]
def get_training_inference_pipeline(args):

    # Get the configuration for the pipeline
    with open(args.model_config) as fp:
        model_config = yaml.safe_load(fp)
    # Add it to the W&B configuration so the values for the hyperparams
    # are tracked
    wandb.config.update(model_config)

    # 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(model_config["features"]["categorical"])
    categorical_transformer = make_pipeline(
        SimpleImputer(strategy="constant", fill_value=0), OrdinalEncoder()
    )
    # Numerical preprocessing pipeline
    numeric_features = sorted(model_config["features"]["numerical"])
    numeric_transformer = make_pipeline(
        SimpleImputer(strategy="median"), StandardScaler()
    )
    # Textual ("nlp") preprocessing pipeline
    nlp_features = sorted(model_config["features"]["nlp"])
    # 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, max_features=model_config["tfidf"]["max_features"]
        ),
    )
    # 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
    )

    # Append classifier to preprocessing pipeline.
    # Now we have a full prediction pipeline.
    pipe = Pipeline(
        steps=[
            ("preprocessor", preprocessor),
            ("classifier", RandomForestClassifier(**model_config["random_forest"])),
        ]
    )
    return pipe
//  [python]
continus_cols = ["shop_duration", "recent", "monetary", "max_amount", "items_count", 
                 "valid_points_sum", "member_day", "frequence", "avg_amount", "item_count_turn", 
                 "avg_piece_amount", "monetary3","max_amount3", "items_count3", 
                 "frequence3", "shops_count", "promote_percent", "wxapp_diff", "store_diff", 
                 "week_percent"]

category_cols = ["CHANNEL_NUM_ID", "shop_channel", "infant_group", "water_product_group", 
                 "meat_group", "beauty_group", "health_group", "fruits_group", "vegetables_group", 
                 "pets_group", "snacks_group", "smoke_group", "milk_group", "instant_group", 
                 "grain_group"]


if __name__ == "__main__":
    df = pd.read_csv("./data/churn.csv")
    train, test = train_test_split(df, test_size=0.2)
    
    # 整体组装
    preprocessor = ColumnTransformer(
            transformers=[
                    ("0_fillna", Imputer(strategy='constant', fill_value=0), continus_cols), 
                    ("onehot", OneHotEncoder(handle_unknown='ignore'), category_cols)
                    ])
    clf = Pipeline(steps=[("preprocessor", preprocessor), 
                          ("classifier", XGBClassifier(max_depth=8, n_estimators=100, n_jods=3))])
    
    # 训练
    clf.fit(train[continus_cols + category_cols], train["label"])
    pickle.dump(clf, open("./churn_xgb.model", "wb"))
    
    # 预测
    model = pickle.load(open("./churn_xgb.model", "rb"))
    predictions = model.predict(test[continus_cols + category_cols])
    predict_proba = model.predict_proba(test[continus_cols + category_cols])[:, 1]
    
    # 模型评价
    print("acc:", accuracy_score(test["label"], predictions))
    print("pri:", precision_score(test["label"], predictions))
    print("rec:", recall_score(test["label"], predictions))
    print("auc:", roc_auc_score(test["label"], predict_proba))
//  [python]
continus_cols = ["shop_duration", "recent", "monetary", "max_amount", "items_count", 
                 "valid_points_sum", "member_day", "frequence", "avg_amount", "item_count_turn", 
                 "avg_piece_amount", "monetary3","max_amount3", "items_count3", 
                 "frequence3", "shops_count", "promote_percent", "wxapp_diff", "store_diff", 
                 "week_percent"]

category_cols = ["CHANNEL_NUM_ID", "shop_channel", "infant_group", "water_product_group", 
                 "meat_group", "beauty_group", "health_group", "fruits_group", "vegetables_group", 
                 "pets_group", "snacks_group", "smoke_group", "milk_group", "instant_group", 
                 "grain_group"]


if __name__ == "__main__":
    df = pd.read_csv("./data/churn.csv")
    train, test = train_test_split(df, test_size=0.2)
    
    # 整体组装
    preprocessor = ColumnTransformer(
            transformers=[
                    ("0_fillna", SimpleImputer(strategy='constant', fill_value=0), continus_cols), 
                    ("onehot", OneHotEncoder(handle_unknown='ignore'), category_cols)
                    ])
    clf = Pipeline(steps=[("preprocessor", preprocessor), 
                          ("classifier", XGBClassifier(max_depth=8, n_estimators=100, n_jods=3))])
    
    # 训练
    clf.fit(train[continus_cols + category_cols], train["label"])
    pickle.dump(clf, open("./churn_xgb.model", "wb"))
    
    # 预测
    model = pickle.load(open("./churn_xgb.model", "rb"))
    predictions = model.predict(test[continus_cols + category_cols])
    predict_proba = model.predict_proba(test[continus_cols + category_cols])[:, 1]
    
    # 模型评价
    print("acc:", accuracy_score(test["label"], predictions))
    print("pri:", precision_score(test["label"], predictions))
    print("rec:", recall_score(test["label"], predictions))
    print("auc:", roc_auc_score(test["label"], predict_proba))
//  [python]
def test_imputation_mean_median_error_invalid_type(strategy, dtype):
    X = np.array([["a", "b", 3],
                  [4, "e", 6],
                  ["g", "h", 9]], dtype=dtype)
    msg = "non-numeric data:\ncould not convert string to float: '"
    with pytest.raises(ValueError, match=msg):
        imputer = Imputer(strategy=strategy)
        imputer.fit_transform(X)
//  [python]
def test_imputation_mean_median_error_invalid_type(strategy, dtype):
    X = np.array([["a", "b", 3],
                  [4, "e", 6],
                  ["g", "h", 9]], dtype=dtype)
    msg = "non-numeric data:\ncould not convert string to float: '"
    with pytest.raises(ValueError, match=msg):
        imputer = SimpleImputer(strategy=strategy)
        imputer.fit_transform(X)
//  [python]
def test_heterogeneous_ensemble_support_missing_values(Ensemble,
                                                       Estimator, X, y):
    # check that Voting and Stacking predictor delegate the missing values
    # validation to the underlying estimator.
    X = X.copy()
    mask = np.random.choice([1, 0], X.shape, p=[.1, .9]).astype(bool)
    X[mask] = np.nan
    pipe = make_pipeline(Imputer(), Estimator())
    ensemble = Ensemble(estimators=[('pipe1', pipe), ('pipe2', pipe)])
    ensemble.fit(X, y).score(X, y)
//  [python]
def test_heterogeneous_ensemble_support_missing_values(Ensemble,
                                                       Estimator, X, y):
    # check that Voting and Stacking predictor delegate the missing values
    # validation to the underlying estimator.
    X = X.copy()
    mask = np.random.choice([1, 0], X.shape, p=[.1, .9]).astype(bool)
    X[mask] = np.nan
    pipe = make_pipeline(SimpleImputer(), Estimator())
    ensemble = Ensemble(estimators=[('pipe1', pipe), ('pipe2', pipe)])
    ensemble.fit(X, y).score(X, y)
//  [python]
def test_imputation_mean_median_error_invalid_type_list_pandas(strategy, type):
    X = [["a", "b", 3],
         [4, "e", 6],
         ["g", "h", 9]]
    if type == 'dataframe':
        pd = pytest.importorskip("pandas")
        X = pd.DataFrame(X)
    msg = "non-numeric data:\ncould not convert string to float: '"
    with pytest.raises(ValueError, match=msg):
        imputer = Imputer(strategy=strategy)
        imputer.fit_transform(X)
//  [python]
def test_imputation_mean_median_error_invalid_type_list_pandas(strategy, type):
    X = [["a", "b", 3],
         [4, "e", 6],
         ["g", "h", 9]]
    if type == 'dataframe':
        pd = pytest.importorskip("pandas")
        X = pd.DataFrame(X)
    msg = "non-numeric data:\ncould not convert string to float: '"
    with pytest.raises(ValueError, match=msg):
        imputer = SimpleImputer(strategy=strategy)
        imputer.fit_transform(X)
//  [python]
def test_imputation_pipeline_grid_search():
    # Test imputation within a pipeline + gridsearch.
    X = _sparse_random_matrix(100, 100, density=0.10)
    missing_values = X.data[0]

    pipeline = Pipeline([('imputer',
                          Imputer(missing_values=missing_values)),
                         ('tree',
                          tree.DecisionTreeRegressor(random_state=0))])

    parameters = {
        'imputer__strategy': ["mean", "median", "most_frequent"]
    }

    Y = _sparse_random_matrix(100, 1, density=0.10).toarray()
    gs = GridSearchCV(pipeline, parameters)
    gs.fit(X, Y)
//  [python]
def test_imputation_pipeline_grid_search():
    # Test imputation within a pipeline + gridsearch.
    X = _sparse_random_matrix(100, 100, density=0.10)
    missing_values = X.data[0]

    pipeline = Pipeline([('imputer',
                          SimpleImputer(missing_values=missing_values)),
                         ('tree',
                          tree.DecisionTreeRegressor(random_state=0))])

    parameters = {
        'imputer__strategy': ["mean", "median", "most_frequent"]
    }

    Y = _sparse_random_matrix(100, 1, density=0.10).toarray()
    gs = GridSearchCV(pipeline, parameters)
    gs.fit(X, Y)
//  [python]
def test_simple_imputation_string_list(strategy, expected):
    X = [['a', 'b'],
         ['c', np.nan]]

    X_true = np.array([
        ['a', 'b'],
        ['c', expected]
    ], dtype=object)

    imputer = Imputer(strategy=strategy)
    X_trans = imputer.fit_transform(X)

    assert_array_equal(X_trans, X_true)
//  [python]
def test_simple_imputation_string_list(strategy, expected):
    X = [['a', 'b'],
         ['c', np.nan]]

    X_true = np.array([
        ['a', 'b'],
        ['c', expected]
    ], dtype=object)

    imputer = SimpleImputer(strategy=strategy)
    X_trans = imputer.fit_transform(X)

    assert_array_equal(X_trans, X_true)
//  [python]
def __init__(self, params: Optional[OperationParameters] = None):
        super().__init__(params)
        default_params_categorical = {'strategy': 'most_frequent'}
        self.params_cat = {**self.params.to_dict(), **default_params_categorical}
        self.params_num = self.params.to_dict()
        self.categorical_ids = None
        self.non_categorical_ids = None
        self.ids_binary_integer_features = {}

        self.imputer_cat = sklearn.preprocessing.Imputer(**self.params_cat)
        self.imputer_num = sklearn.preprocessing.Imputer(**self.params_num)
//  [python]
def __init__(self, params: Optional[OperationParameters] = None):
        super().__init__(params)
        default_params_categorical = {'strategy': 'most_frequent'}
        self.params_cat = {**self.params.to_dict(), **default_params_categorical}
        self.params_num = self.params.to_dict()
        self.categorical_ids = None
        self.non_categorical_ids = None
        self.ids_binary_integer_features = {}

        self.imputer_cat = SimpleImputer(**self.params_cat)
        self.imputer_num = SimpleImputer(**self.params_num)
//  [python]
def test_simple_imputation_inverse_transform_exceptions(missing_value):
    X_1 = np.array([
        [9, missing_value, 3, -1],
        [4, -1, 5, 4],
        [6, 7, missing_value, -1],
        [8, 9, 0, missing_value]
    ])

    imputer = Imputer(missing_values=missing_value, strategy="mean")
    X_1_trans = imputer.fit_transform(X_1)
    with pytest.raises(ValueError,
                       match=f"Got 'add_indicator={imputer.add_indicator}'"):
        imputer.inverse_transform(X_1_trans)
//  [python]
def test_simple_imputation_inverse_transform_exceptions(missing_value):
    X_1 = np.array([
        [9, missing_value, 3, -1],
        [4, -1, 5, 4],
        [6, 7, missing_value, -1],
        [8, 9, 0, missing_value]
    ])

    imputer = SimpleImputer(missing_values=missing_value, strategy="mean")
    X_1_trans = imputer.fit_transform(X_1)
    with pytest.raises(ValueError,
                       match=f"Got 'add_indicator={imputer.add_indicator}'"):
        imputer.inverse_transform(X_1_trans)
//  [python]
def test_feature_names_out(self):
        for encoder_name in encoders.__all__:
            if sklearn.__version__ < "1.2.0":
                continue
            else:
                sklearn.set_config(transform_output="pandas")
            with self.subTest(encoder_name=encoder_name):
                encoder = getattr(encoders, encoder_name)()
                X_t = encoder.fit_transform(X, y)

                categorical_preprocessor_start = Pipeline(
                    steps=[
                        ("encoder", getattr(encoders, encoder_name)())
                    ]
                )
                categorical_preprocessor_middle = Pipeline(
                    steps=[
                        ("imputation_constant", Imputer(fill_value="missing", strategy="constant")),
                        ("encoder", getattr(encoders, encoder_name)())
                    ]
                )
                numerical_preprocessor = Pipeline(
                    steps=[
                        ("imputation_constant", Imputer(fill_value=0, strategy="constant"))
                    ]
                )
                preprocessor = ColumnTransformer(
                    [
                        ("categorical_prep_start", categorical_preprocessor_start, ["categorical", "na_categorical"]),
                        ("categorical_prep_middle", categorical_preprocessor_middle, ["categorical", "na_categorical"]),
                        ("numerical_prep", numerical_preprocessor, ["float"])
                    ]
                )
                X_tt = preprocessor.fit_transform(X, y)

                assert_array_equal(
                    np.array(X_t.columns),
                    encoder.get_feature_names_out()
                )
                assert_array_equal(
                    np.array(X_tt.columns),
                    preprocessor.get_feature_names_out()
                )
                assert_array_equal(
                    np.array(
                        [
                            c
                            for c in X_t.columns
                            if c not in num_columns
                        ]
                    ),
                    np.array(
                        [
                            c[len("categorical_prep_start__"):]
                            for c in X_tt.columns
                            if "categorical_prep_start" in c
                        ]
                    )
                )
                assert_array_equal(
                    np.array(
                        [
                            c
                            for c in X_t.columns
                            if c not in num_columns
                        ]
                    ),
                    np.array(
                        [
                            c[len("categorical_prep_middle__"):]
                            for c in X_tt.columns
                            if "categorical_prep_middle" in c
                        ]
                    )
                )
            sklearn.set_config(transform_output="default")
//  [python]
def test_feature_names_out(self):
        for encoder_name in encoders.__all__:
            if sklearn.__version__ < "1.2.0":
                continue
            else:
                sklearn.set_config(transform_output="pandas")
            with self.subTest(encoder_name=encoder_name):
                encoder = getattr(encoders, encoder_name)()
                X_t = encoder.fit_transform(X, y)

                categorical_preprocessor_start = Pipeline(
                    steps=[
                        ("encoder", getattr(encoders, encoder_name)())
                    ]
                )
                categorical_preprocessor_middle = Pipeline(
                    steps=[
                        ("imputation_constant", SimpleImputer(fill_value="missing", strategy="constant")),
                        ("encoder", getattr(encoders, encoder_name)())
                    ]
                )
                numerical_preprocessor = Pipeline(
                    steps=[
                        ("imputation_constant", SimpleImputer(fill_value=0, strategy="constant"))
                    ]
                )
                preprocessor = ColumnTransformer(
                    [
                        ("categorical_prep_start", categorical_preprocessor_start, ["categorical", "na_categorical"]),
                        ("categorical_prep_middle", categorical_preprocessor_middle, ["categorical", "na_categorical"]),
                        ("numerical_prep", numerical_preprocessor, ["float"])
                    ]
                )
                X_tt = preprocessor.fit_transform(X, y)

                assert_array_equal(
                    np.array(X_t.columns),
                    encoder.get_feature_names_out()
                )
                assert_array_equal(
                    np.array(X_tt.columns),
                    preprocessor.get_feature_names_out()
                )
                assert_array_equal(
                    np.array(
                        [
                            c
                            for c in X_t.columns
                            if c not in num_columns
                        ]
                    ),
                    np.array(
                        [
                            c[len("categorical_prep_start__"):]
                            for c in X_tt.columns
                            if "categorical_prep_start" in c
                        ]
                    )
                )
                assert_array_equal(
                    np.array(
                        [
                            c
                            for c in X_t.columns
                            if c not in num_columns
                        ]
                    ),
                    np.array(
                        [
                            c[len("categorical_prep_middle__"):]
                            for c in X_tt.columns
                            if "categorical_prep_middle" in c
                        ]
                    )
                )
            sklearn.set_config(transform_output="default")
//  [python]
numeric_preprocessor = Pipeline(
    steps=[
        ("imputation_mean", Imputer(missing_values=np.nan, strategy="mean")),
        ("scaler", StandardScaler()),
    ]
)

categorical_preprocessor = Pipeline(
    steps=[
        (
            "imputation_constant",
            Imputer(fill_value="missing", strategy="constant"),
        ),
        ("onehot", OneHotEncoder(handle_unknown="ignore")),
    ]
)
//  [python]
numeric_preprocessor = Pipeline(
    steps=[
        ("imputation_mean", SimpleImputer(missing_values=np.nan, strategy="mean")),
        ("scaler", StandardScaler()),
    ]
)

categorical_preprocessor = Pipeline(
    steps=[
        (
            "imputation_constant",
            SimpleImputer(fill_value="missing", strategy="constant"),
        ),
        ("onehot", OneHotEncoder(handle_unknown="ignore")),
    ]
)
//  [python]
def test_data_reconstruction_drift_calculator_given_wrong_cat_imputer_object_raises_typeerror(  # noqa: D103
    sample_drift_data_with_nans,
):
    with pytest.raises(TypeError):
        DataReconstructionDriftCalculator(
            column_names=['f1', 'f2', 'f3', 'f4'],
            timestamp_column_name='timestamp',
            chunk_period='W',
            imputer_categorical=5,
            imputer_continuous=Imputer(missing_values=np.nan, strategy='mean'),
        )
//  [python]
def test_data_reconstruction_drift_calculator_given_wrong_cat_imputer_object_raises_typeerror(  # noqa: D103
    sample_drift_data_with_nans,
):
    with pytest.raises(TypeError):
        DataReconstructionDriftCalculator(
            column_names=['f1', 'f2', 'f3', 'f4'],
            timestamp_column_name='timestamp',
            chunk_period='W',
            imputer_categorical=5,
            imputer_continuous=SimpleImputer(missing_values=np.nan, strategy='mean'),
        )