# Buggy code
def munge_augment(
    X,
    feature_types_metadata: FeatureTypesMetadata,
    num_augmented_samples=10000,
    perturb_prob=0.5,
    s=1.0,
    **kwargs,
):
    """Uses MUNGE algorithm to generate synthetic datapoints for learning to mimic teacher model in distillation: https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf
    Args:
        num_augmented_samples: number of additional augmented data points to return
        perturb_prob: probability of perturbing each feature during augmentation. Set near 0 to ensure augmented sample distribution remains closer to real data.
        s: We noise numeric features by their std-dev divided by this factor (inverse of continuous_feature_noise). Set large to ensure augmented sample distribution remains closer to real data.
    """
    nn_dummy = TabularNeuralNetModel(
        path="nn_dummy",
        name="nn_dummy",
        problem_type=REGRESSION,
        eval_metric=mean_squared_error,
        hyperparameters={
            "num_dataloading_workers": 0,
            "proc.embed_min_categories": np.inf,
        },
        features=list(X.columns),
        feature_types_metadata=feature_types_metadata,
    )
    processed_data = nn_dummy.process_train_data(
        df=nn_dummy.preprocess(X),
        labels=pd.Series([1] * len(X)),
        batch_size=nn_dummy.params["batch_size"],
        num_dataloading_workers=0,
        impute_strategy=nn_dummy.params["proc.impute_strategy"],
        max_category_levels=nn_dummy.params["proc.max_category_levels"],
        skew_threshold=nn_dummy.params["proc.skew_threshold"],
        embed_min_categories=nn_dummy.params["proc.embed_min_categories"],
        use_ngram_features=nn_dummy.params["use_ngram_features"],
    )
    X_vector = processed_data.dataset._data[processed_data.vectordata_index].asnumpy()
    processed_data = None
    nn_dummy = None
    gc.collect()

    neighbor_finder = NearestNeighbors(n_neighbors=2)
    neighbor_finder.fit(X_vector)
    neigh_dist, neigh_ind = neighbor_finder.kneighbors(X_vector)
    neigh_ind = neigh_ind[:, 1]  # contains indices of nearest neighbors
    neigh_dist = None
    # neigh_dist = neigh_dist[:,1]  # contains distances to nearest neighbors
    neighbor_finder = None
    gc.collect()

    if perturb_prob > 1.0:
        raise ValueError("frac_perturb must be <= 1")
    logger.log(
        20,
        f"MUNGE: Augmenting training data with {num_augmented_samples} synthetic samples for distillation...",
    )
    X = X.copy()
    X_aug = pd.concat([X.iloc[[0]].copy()] * num_augmented_samples)
    X_aug.reset_index(drop=True, inplace=True)
    continuous_types = ["float", "int"]
    continuous_featnames = []  # these features will have shuffled values with added noise
    for contype in continuous_types:
        if contype in feature_types_metadata.feature_types_raw:
            continuous_featnames += feature_types_metadata.feature_types_raw[contype]
    for col in continuous_featnames:
        X_aug[col] = X_aug[col].astype(float)
        X[col] = X[col].astype(float)

    for i in range(num_augmented_samples):
        og_ind = i % len(X)
        augdata_i = X.iloc[og_ind].copy()
        neighbor_i = X.iloc[neigh_ind[og_ind]].copy()
        # dist_i = neigh_dist[og_ind]
        cols_toperturb = np.random.choice(
            list(X.columns),
            size=np.random.binomial(X.shape[1], p=perturb_prob, size=1)[0],
            replace=False,
        )
        for col in cols_toperturb:
            new_val = neighbor_i[col]
            if col in continuous_featnames:
                new_val += np.random.normal(scale=np.abs(augdata_i[col] - new_val) / s)
            augdata_i[col] = new_val
        X_aug.iloc[i] = augdata_i

    return X_aug


# Traceback (ValueError)
File delimiter for ../input/house-prices-advanced-regression-techniques/test.csv inferred as ',' (comma). If this is incorrect, please manually load the data as a pandas DataFrame.
Loaded data from: ../input/house-prices-advanced-regression-techniques/test.csv | Columns = 80 / 80 | Rows = 1459 -> 1459
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-48-91455d7b9940> in <module>
1 test_data = task.Dataset(file_path=os.path.join(in_dir, job, 'test.csv'))
2
----> 3 preds = model.predict(test_data)

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/task/tabular_prediction/predictor.py in predict(self, dataset, model, as_pandas, use_pred_cache, add_to_pred_cache)
130         """
131         dataset = self.__get_dataset(dataset)
--> 132         return self._learner.predict(X=dataset, model=model, as_pandas=as_pandas, use_pred_cache=use_pred_cache, add_to_pred_cache=add_to_pred_cache)
133
134     def predict_proba(self, dataset, model=None, as_pandas=False, as_multiclass=False):

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/utils/tabular/ml/learner/abstract_learner.py in predict(self, X, model, as_pandas, use_pred_cache, add_to_pred_cache)
127
128         if len(X_cache_miss) > 0:
--> 129             y_pred_proba = self.predict_proba(X=X_cache_miss, model=model, inverse_transform=False)
130             problem_type = self.trainer_problem_type or self.problem_type
131             y_pred = get_pred_from_proba(y_pred_proba=y_pred_proba, problem_type=problem_type)

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/utils/tabular/ml/learner/abstract_learner.py in predict_proba(self, X, model, as_pandas, as_multiclass, inverse_transform)
92     # TODO: Add pred_proba_cache functionality as in predict()
93     def predict_proba(self, X: DataFrame, model=None, as_pandas=False, as_multiclass=False, inverse_transform=True):
---> 94         X = self.transform_features(X)
95         y_pred_proba = self.load_trainer().predict_proba(X, model=model)
96         if inverse_transform:

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/utils/tabular/ml/learner/abstract_learner.py in transform_features(self, X)
216     def transform_features(self, X):
217         for feature_generator in self.feature_generators:
--> 218             X = feature_generator.transform(X)
219         return X
220

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/utils/tabular/utils/decorators.py in inner1(*args, **kwargs)
12         begin = time.time()
13
---> 14         output = func(*args, **kwargs)
15
16         # storing time after function execution

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/autogluon/utils/tabular/features/abstract_feature_generator.py in transform(self, X)
183             raise ValueError(f'Required columns are missing from the provided dataset. Missing columns: {missing_cols}')
184
--> 185         X = X.astype(self.features_init_types)
186         X.reset_index(drop=True, inplace=True)
187         X_features = self.generate_features(X)

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/generic.py in astype(self, dtype, copy, errors, **kwargs)
5863                     results.append(
5864                         col.astype(
-> 5865                             dtype=dtype[col_name], copy=copy, errors=errors, **kwargs
5866                         )
5867                     )

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/generic.py in astype(self, dtype, copy, errors, **kwargs)
5880             # else, only a single dtype is given
5881             new_data = self._data.astype(
-> 5882                 dtype=dtype, copy=copy, errors=errors, **kwargs
5883             )
5884             return self._constructor(new_data).__finalize__(self)

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/internals/managers.py in astype(self, dtype, **kwargs)
579
580     def astype(self, dtype, **kwargs):
--> 581         return self.apply("astype", dtype=dtype, **kwargs)
582
583     def convert(self, **kwargs):

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/internals/managers.py in apply(self, f, axes, filter, do_integrity_check, consolidate, **kwargs)
436                     kwargs[k] = obj.reindex(b_items, axis=axis, copy=align_copy)
437
--> 438             applied = getattr(b, f)(**kwargs)
439             result_blocks = _extend_blocks(applied, result_blocks)
440

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/internals/blocks.py in astype(self, dtype, copy, errors, values, **kwargs)
557
558     def astype(self, dtype, copy=False, errors="raise", values=None, **kwargs):
--> 559         return self._astype(dtype, copy=copy, errors=errors, values=values, **kwargs)
560
561     def _astype(self, dtype, copy=False, errors="raise", values=None, **kwargs):

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/internals/blocks.py in _astype(self, dtype, copy, errors, values, **kwargs)
641                     # _astype_nansafe works fine with 1-d only
642                     vals1d = values.ravel()
--> 643                     values = astype_nansafe(vals1d, dtype, copy=True, **kwargs)
644
645                 # TODO(extension)

~/miniconda3/envs/autogluon/lib/python3.7/site-packages/pandas/core/dtypes/cast.py in astype_nansafe(arr, dtype, copy, skipna)
698         if not np.isfinite(arr).all():
699             raise ValueError(
--> 700                 "Cannot convert non-finite values (NA or inf) to " "integer"
701             )
702

ValueError: Cannot convert non-finite values (NA or inf) to integer

# Fixed code
def munge_augment(
    X,
    feature_metadata: FeatureMetadata,
    num_augmented_samples=10000,
    perturb_prob=0.5,
    s=1.0,
    **kwargs,
):
    """Uses MUNGE algorithm to generate synthetic datapoints for learning to mimic teacher model in distillation: https://www.cs.cornell.edu/~caruana/compression.kdd06.pdf
    Args:
        num_augmented_samples: number of additional augmented data points to return
        perturb_prob: probability of perturbing each feature during augmentation. Set near 0 to ensure augmented sample distribution remains closer to real data.
        s: We noise numeric features by their std-dev divided by this factor (inverse of continuous_feature_noise). Set large to ensure augmented sample distribution remains closer to real data.
    """
    nn_dummy = TabularNeuralNetModel(
        path="nn_dummy",
        name="nn_dummy",
        problem_type=REGRESSION,
        eval_metric=mean_squared_error,
        hyperparameters={
            "num_dataloading_workers": 0,
            "proc.embed_min_categories": np.inf,
        },
        features=list(X.columns),
        feature_metadata=feature_metadata,
    )
    processed_data = nn_dummy.process_train_data(
        df=nn_dummy.preprocess(X),
        labels=pd.Series([1] * len(X)),
        batch_size=nn_dummy.params["batch_size"],
        num_dataloading_workers=0,
        impute_strategy=nn_dummy.params["proc.impute_strategy"],
        max_category_levels=nn_dummy.params["proc.max_category_levels"],
        skew_threshold=nn_dummy.params["proc.skew_threshold"],
        embed_min_categories=nn_dummy.params["proc.embed_min_categories"],
        use_ngram_features=nn_dummy.params["use_ngram_features"],
    )
    X_vector = processed_data.dataset._data[processed_data.vectordata_index].asnumpy()
    processed_data = None
    nn_dummy = None
    gc.collect()

    neighbor_finder = NearestNeighbors(n_neighbors=2)
    neighbor_finder.fit(X_vector)
    neigh_dist, neigh_ind = neighbor_finder.kneighbors(X_vector)
    neigh_ind = neigh_ind[:, 1]  # contains indices of nearest neighbors
    neigh_dist = None
    # neigh_dist = neigh_dist[:,1]  # contains distances to nearest neighbors
    neighbor_finder = None
    gc.collect()

    if perturb_prob > 1.0:
        raise ValueError("frac_perturb must be <= 1")
    logger.log(
        20,
        f"MUNGE: Augmenting training data with {num_augmented_samples} synthetic samples for distillation...",
    )
    X = X.copy()
    X_aug = pd.concat([X.iloc[[0]].copy()] * num_augmented_samples)
    X_aug.reset_index(drop=True, inplace=True)
    continuous_types = ["float", "int"]
    continuous_featnames = feature_metadata.get_features(
        valid_raw_types=continuous_types
    )  # these features will have shuffled values with added noise
    for col in continuous_featnames:
        X_aug[col] = X_aug[col].astype(float)
        X[col] = X[col].astype(float)

    for i in range(num_augmented_samples):
        og_ind = i % len(X)
        augdata_i = X.iloc[og_ind].copy()
        neighbor_i = X.iloc[neigh_ind[og_ind]].copy()
        # dist_i = neigh_dist[og_ind]
        cols_toperturb = np.random.choice(
            list(X.columns),
            size=np.random.binomial(X.shape[1], p=perturb_prob, size=1)[0],
            replace=False,
        )
        for col in cols_toperturb:
            new_val = neighbor_i[col]
            if col in continuous_featnames:
                new_val += np.random.normal(scale=np.abs(augdata_i[col] - new_val) / s)
            augdata_i[col] = new_val
        X_aug.iloc[i] = augdata_i

    return X_aug