Skip to content

Commit

Permalink
fix: check target is not None in converting (#1316)
Browse files Browse the repository at this point in the history
  • Loading branch information
Lopa10ko authored Aug 1, 2024
1 parent 0368e9c commit 3d49382
Showing 1 changed file with 12 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,12 @@ def plot_feature_importance(self, importance_type='weight'):
@staticmethod
def convert_to_dataframe(data: Optional[InputData], identify_cats: bool):
dataframe = pd.DataFrame(data=data.features)
dataframe['target'] = np.ravel(data.target)
if data.target is not None:
dataframe['target'] = np.ravel(data.target)
else:
# TODO: temp workaround in case data.target is set to None intentionally
# for test.integration.models.test_model.check_predict_correct
dataframe['target'] = np.zeros(len(data.features))

if identify_cats and data.categorical_idx is not None:
for col in dataframe.columns[data.categorical_idx]:
Expand Down Expand Up @@ -231,7 +236,12 @@ def set_eval_metric(n_classes):
@staticmethod
def convert_to_dataframe(data: Optional[InputData], identify_cats: bool):
dataframe = pd.DataFrame(data=data.features, columns=data.features_names)
dataframe['target'] = np.ravel(data.target)
if data.target is not None:
dataframe['target'] = np.ravel(data.target)
else:
# TODO: temp workaround in case data.target is set to None intentionally
# for test.integration.models.test_model.check_predict_correct
dataframe['target'] = np.zeros(len(data.features))

if identify_cats and data.categorical_idx is not None:
for col in dataframe.columns[data.categorical_idx]:
Expand Down

0 comments on commit 3d49382

Please sign in to comment.