Investigate ACTGAN NaN handling bug

Co-authored-by: marjan_emd <marjan_emd> GitOrigin-RevId: 7d20e3b6120289a420d4bbb1529c8820927defda
gretelai · Mar 12, 2024 · 5c9534a · 5c9534a
1 parent a63c68d
commit 5c9534a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/gretel_synthetics/actgan/transformers.py b/src/gretel_synthetics/actgan/transformers.py
@@ -193,7 +193,7 @@ def _fit(self, data: SeriesOrDFLike) -> None:
             data = pd.Series(data)
 
         if self.handle_rounding_nan == MODE:
-            self._mode_values = frozenset(list(data.mode()))
+            self._mode_values = frozenset(list(data.mode(dropna=False)))
         self._nan_proxy = _new_uuid()
         data = data.fillna(self._nan_proxy)
 
@@ -229,6 +229,7 @@ def _transform(self, data: ListOrSeriesOrDF) -> np.ndarray:
         #         'please fit the transformer again with the new data.'
         #     )
 
+        data = data.fillna(self._nan_proxy)
         ndarray = self.encoder.transform(data).to_numpy()
         return ndarray