From 5edb334a0e77f44f0b3b29c6de8e63914760a85d Mon Sep 17 00:00:00 2001 From: Christopher Canel Date: Thu, 8 Feb 2024 14:43:45 +0000 Subject: [PATCH] Do not replace unknowns in HGBDT model --- unfair/model/data.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unfair/model/data.py b/unfair/model/data.py index f468178..11fbfc0 100644 --- a/unfair/model/data.py +++ b/unfair/model/data.py @@ -216,7 +216,9 @@ def extract_fets(dat, split_name, net): or np.isinf(dat_out[features.LABEL_FET]).any() ), f'Warning: NaNs or Infs in ground truth for split "{split_name}".' - replace_unknowns(dat_in, is_dt) + # HGBDT can handle unknowns, but other model types cannot. + if not isinstance(net, models.HistGbdtSklearnWrapper): + replace_unknowns(dat_in, is_dt) # Convert output features to class labels. dat_out = net.convert_to_class(dat_out)