refactor the method '_load_source' of the class TrainConfig

tdspora · Jan 9, 2025 · 3e6d68f · 3e6d68f
1 parent 83c3155
commit 3e6d68f
Showing 1 changed file with 9 additions and 9 deletions.
diff --git a/src/syngen/ml/config/configurations.py b/src/syngen/ml/config/configurations.py
@@ -118,18 +118,18 @@ def _load_source(self) -> Tuple[pd.DataFrame, Dict]:
         if self.loader is not None:
             return self._fetch_dataframe()
         if os.path.exists(self.paths["path_to_flatten_metadata"]):
-            data, schema = DataLoader(self.paths["input_data_path"]).load_data()
-            self.original_schema = DataLoader(self.paths["input_data_path"]).original_schema
+            data_loader = DataLoader(self.paths["input_data_path"])
+            data, _ = data_loader.load_data()
             schema = CSVConvertor.schema
-            return data, schema
         else:
             data_loader = DataLoader(self.source)
-            self.original_schema = data_loader.original_schema
-            if self.original_schema is not None:
-                logger.trace(
-                    f"The schema of the table - '{self.table_name}': {self.original_schema}"
-                )
-            return data_loader.load_data()
+            data, schema = data_loader.load_data()
+        self.original_schema = data_loader.original_schema
+        if self.original_schema is not None:
+            logger.trace(
+                f"The schema of the table - '{self.table_name}': {self.original_schema}"
+            )
+        return data, schema
 
     def _remove_empty_columns(self):
         """