fix: inf filling

Signed-off-by: Avik Basu <ab93@users.noreply.github.com>
numaproj · Nov 21, 2023 · fb80cb6 · fb80cb6
1 parent 5aa4c13
commit fb80cb6
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/numalogic/udfs/trainer/_base.py b/numalogic/udfs/trainer/_base.py
@@ -256,6 +256,7 @@ def _is_data_sufficient(self, payload: TrainerPayload, df: pd.DataFrame) -> bool
             return False
         return True
 
+    # TODO: Use a custom imputer in transforms module
     @staticmethod
     def get_feature_arr(
         raw_df: pd.DataFrame, metrics: list[str], fill_value: float = 0.0
@@ -265,7 +266,7 @@ def get_feature_arr(
             if col not in raw_df.columns:
                 raw_df[col] = fill_value
         feat_df = raw_df[metrics]
-        feat_df = feat_df.fillna(fill_value)
+        feat_df = feat_df.fillna(fill_value).replace([np.inf, -np.inf], fill_value)
         return feat_df.to_numpy(dtype=np.float32)
 
     def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.6.1.dev3"
+version = "0.6.1.dev4"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]