refactor the code in the class CSVLoader, the class BivariateMetric

tdspora · Nov 6, 2024 · 6c8a689 · 6c8a689
1 parent f98c017
commit 6c8a689
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/syngen/ml/data_loaders/data_loaders.py b/src/syngen/ml/data_loaders/data_loaders.py
@@ -257,7 +257,7 @@ def _save_data(self, df: pd.DataFrame, **kwargs):
             ):
                 filtered_kwargs["na_rep"] = format_params["na_values"][0]
                 logger.warning(
-                    "Since the 'na_values' parameter is not empty, "
+                    "Since the 'na_values' parameter in the 'format' sections is not empty, "
                     "the missing values will be filled with "
                     "the first value from the 'na_values' parameter"
                 )

diff --git a/src/syngen/ml/metrics/metrics_classes/metrics.py b/src/syngen/ml/metrics/metrics_classes/metrics.py
@@ -352,7 +352,7 @@ def __init__(
         nan_labels_dict = self.dataset.nan_labels_dict
         na_values = self.dataset.format.get("na_values", [])
         self.missing_values: Dict[str, str] = (
-            {col: na_values[0] for col in self.dataset.order_of_columns}
+            {col: nan_labels_dict.get(col, na_values[0]) for col in self.dataset.order_of_columns}
             if na_values
             else nan_labels_dict
         )