dragon-ml-toolbox 13.6.0__py3-none-any.whl → 13.8.0__py3-none-any.whl

{dragon_ml_toolbox-13.6.0.dist-info → dragon_ml_toolbox-13.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragon-ml-toolbox
-Version: 13.6.0
+Version: 13.8.0
 Summary: A collection of tools for data science and machine learning projects.
 Author-email: "Karl L. Loza Vidaurre" <luigiloza@gmail.com>
 License-Expression: MIT

{dragon_ml_toolbox-13.6.0.dist-info → dragon_ml_toolbox-13.8.0.dist-info}/RECORD

@@ -1,9 +1,9 @@
-dragon_ml_toolbox-13.6.0.dist-info/licenses/LICENSE,sha256=L35WDmmLZNTlJvxF6Vy7Uy4SYNi6rCfWUqlTHpoRMoU,1081
-dragon_ml_toolbox-13.6.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=iy2r_R7wjzsCbz_Q_jMsp_jfZ6oP8XW9QhwzRBH0mGY,1904
+dragon_ml_toolbox-13.8.0.dist-info/licenses/LICENSE,sha256=L35WDmmLZNTlJvxF6Vy7Uy4SYNi6rCfWUqlTHpoRMoU,1081
+dragon_ml_toolbox-13.8.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=iy2r_R7wjzsCbz_Q_jMsp_jfZ6oP8XW9QhwzRBH0mGY,1904
 ml_tools/ETL_cleaning.py,sha256=2VBRllV8F-ZiPylPp8Az2gwn5ztgazN0BH5OKnRUhV0,20402
 ml_tools/ETL_engineering.py,sha256=KfYqgsxupAx6e_TxwO1LZXeu5mFkIhVXJrNjP3CzIZc,54927
 ml_tools/GUI_tools.py,sha256=Va6ig-dHULPVRwQYYtH3fvY5XPIoqRcJpRW8oXC55Hw,45413
-ml_tools/MICE_imputation.py,sha256=X273Qlgoqqg7KTmoKd75YDyAPB0UIbTzGP3xsCmRh3E,11717
+ml_tools/MICE_imputation.py,sha256=KLJXGQLKJ6AuWWttAG-LCCaxpS-ygM4dXPiguHDaL6Y,20815
 ml_tools/ML_callbacks.py,sha256=elD2Yr030sv_6gX_m9GVd6HTyrbmt34nFS8lrgS4HtM,15808
 ml_tools/ML_datasetmaster.py,sha256=6caWbq6eu1RE9V51gmceD71PtMctJRjFuLvkkK5ChiY,36271
 ml_tools/ML_evaluation.py,sha256=li77AuP53pCzgrj6p-jTCNtPFgS9Y9XnMWIZn1ulTBM,18946
@@ -13,7 +13,7 @@ ml_tools/ML_models.py,sha256=UVWJHPLVIvFno_csCHH1FwBfTwQ5nX0V8F1TbOByZ4I,31388
 ml_tools/ML_optimization.py,sha256=P0zkhKAwTpkorIBtR0AOIDcyexo5ngmvFUzo3DfNO-E,22692
 ml_tools/ML_scaler.py,sha256=tw6onj9o8_kk3FQYb930HUzvv1zsFZe2YZJdF3LtHkU,7538
 ml_tools/ML_trainer.py,sha256=ZxeOagXW5adFhYIH-oMTlcrLU6VHe4R1EROI7yypNwQ,29665
-ml_tools/ML_utilities.py,sha256=EnKpPTnJ2qjZmz7kvows4Uu5CfSA7ByRmI1v2-KarKw,9337
+ml_tools/ML_utilities.py,sha256=QC44y5mAzA6iUdb3py0bjI-nPjxUatZTdm8sMrb3He0,19364
 ml_tools/PSO_optimization.py,sha256=T-HWHMRJUnPvPwixdU5jif3_rnnI36TzcL8u3oSCwuA,22960
 ml_tools/RNN_forecast.py,sha256=Qa2KoZfdAvSjZ4yE78N4BFXtr3tTr0Gx7tQJZPotsh0,1967
 ml_tools/SQL.py,sha256=vXLPGfVVg8bfkbBE3HVfyEclVbdJy0TBhuQONtMwSCQ,11234
@@ -23,19 +23,19 @@ ml_tools/_logger.py,sha256=dlp5cGbzooK9YSNSZYB4yjZrOaQUGW8PTrM411AOvL8,4717
 ml_tools/_schema.py,sha256=yu6aWmn_2Z4_AxAtJGDDCIa96y6JcUp-vgnCS013Qmw,3908
 ml_tools/_script_info.py,sha256=21r83LV3RubsNZ_RTEUON6RbDf7Mh4_udweNcvdF_Fk,212
 ml_tools/constants.py,sha256=3br5Rk9cL2IUo638eJuMOGdbGQaWssaUecYEvSeRBLM,3322
-ml_tools/custom_logger.py,sha256=7tSAgRL7e-Ekm7rS1FLDocaPLCnaoKc7VSrtfwCtCEg,10067
+ml_tools/custom_logger.py,sha256=i0cAr1qPnwXDyqQ1itk2o72-2jniRXJNEuST2eW4zF4,11016
 ml_tools/data_exploration.py,sha256=-BbWO7BBFapPi_7ZuWo65VqguJXaBfgFSptrXyoWrDk,51902
 ml_tools/ensemble_evaluation.py,sha256=FGHSe8LBI8_w8LjNeJWOcYQ1UK_mc6fVah8gmSvNVGg,26853
 ml_tools/ensemble_inference.py,sha256=0yLmLNj45RVVoSCLH1ZYJG9IoAhTkWUqEZmLOQTFGTY,9348
 ml_tools/ensemble_learning.py,sha256=vsIED7nlheYI4w2SBzP6SC1AnNeMfn-2A1Gqw5EfxsM,21964
 ml_tools/handle_excel.py,sha256=pfdAPb9ywegFkM9T54bRssDOsX-K7rSeV0RaMz7lEAo,14006
-ml_tools/keys.py,sha256=oykUVLB4Wos3AZomowjtI8AFFC5xnMUH-icNHydRpOk,2275
+ml_tools/keys.py,sha256=CcqE9R9R32osR0vLz0i-3cyv1UlVsDWAHqvlVf8xm_0,2492
 ml_tools/math_utilities.py,sha256=xeKq1quR_3DYLgowcp4Uam_4s3JltUyOnqMOGuAiYWU,8802
 ml_tools/optimization_tools.py,sha256=TYFQ2nSnp7xxs-VyoZISWgnGJghFbsWasHjruegyJRs,12763
 ml_tools/path_manager.py,sha256=CyDU16pOKmC82jPubqJPT6EBt-u-3rGVbxyPIZCvDDY,18432
 ml_tools/serde.py,sha256=c8uDYjYry_VrLvoG4ixqDj5pij88lVn6Tu4NHcPkwDU,6943
-ml_tools/utilities.py,sha256=5iN1-UxR0CVROWOjL0LcQwfuKVDtSColO7vjgwqgZgU,16136
-dragon_ml_toolbox-13.6.0.dist-info/METADATA,sha256=oWfSdro-M75lKcpNj5eaSANwFGl92G9oqFzbN89mrxo,6166
-dragon_ml_toolbox-13.6.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-dragon_ml_toolbox-13.6.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
-dragon_ml_toolbox-13.6.0.dist-info/RECORD,,
+ml_tools/utilities.py,sha256=aWqvYzmxlD74PD5Yqu1VuTekDJeYLQrmPIU_VeVyRp0,22526
+dragon_ml_toolbox-13.8.0.dist-info/METADATA,sha256=mvK0WY75d25CARpUbiDoaK3PHtVgRIEcCauCo7RT6wU,6166
+dragon_ml_toolbox-13.8.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dragon_ml_toolbox-13.8.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
+dragon_ml_toolbox-13.8.0.dist-info/RECORD,,

ml_tools/MICE_imputation.py

@@ -7,19 +7,20 @@ from plotnine import ggplot, labs, theme, element_blank # type: ignore
 from typing import Optional, Union
 
 from .utilities import load_dataframe, merge_dataframes, save_dataframe_filename
-from .math_utilities import threshold_binary_values
+from .math_utilities import threshold_binary_values, discretize_categorical_values
 from .path_manager import sanitize_filename, make_fullpath, list_csv_paths
 from ._logger import _LOGGER
 from ._script_info import _script_info
+from ._schema import FeatureSchema
 
 
 __all__ = [
+    "MiceImputer",
     "apply_mice",
     "save_imputed_datasets",
-    "get_na_column_names",
     "get_convergence_diagnostic",
     "get_imputed_distributions",
-    "run_mice_pipeline"
+    "run_mice_pipeline",
 ]
 
 
@@ -79,7 +80,7 @@ def save_imputed_datasets(save_dir: Union[str, Path], imputed_datasets: list, df
 
 
 #Get names of features that had missing values before imputation
-def get_na_column_names(df: pd.DataFrame):
+def _get_na_column_names(df: pd.DataFrame):
     return [col for col in df.columns if df[col].isna().any()]
 
 
@@ -264,7 +265,7 @@ def run_mice_pipeline(df_path_or_dir: Union[str,Path], target_columns: list[str]
         
         save_imputed_datasets(save_dir=save_datasets_path, imputed_datasets=imputed_datasets, df_targets=df_targets, imputed_dataset_names=imputed_dataset_names)
         
-        imputed_column_names = get_na_column_names(df=df)
+        imputed_column_names = _get_na_column_names(df=df)
         
         get_convergence_diagnostic(kernel=kernel, imputed_dataset_names=imputed_dataset_names, column_names=imputed_column_names, root_dir=save_metrics_path)
         
@@ -278,5 +279,206 @@ def _skip_targets(df: pd.DataFrame, target_cols: list[str]):
     return df_feats, df_targets
 
 
+# modern implementation
+class MiceImputer:
+    """
+    A modern MICE imputation pipeline that uses a FeatureSchema
+    to correctly discretize categorical features after imputation.
+    """
+    def __init__(self, 
+                 schema: FeatureSchema,
+                 iterations: int=20,
+                 resulting_datasets: int = 1,
+                 random_state: int = 101):
+        
+        self.schema = schema
+        self.random_state = random_state
+        self.iterations = iterations
+        self.resulting_datasets = resulting_datasets
+        
+        # --- Store schema info ---
+        
+        # 1. Categorical info
+        if not self.schema.categorical_index_map:
+            _LOGGER.warning("FeatureSchema has no 'categorical_index_map'. No discretization will be applied.")
+            self.cat_info = {}
+        else:
+            self.cat_info = self.schema.categorical_index_map
+            
+        # 2. Ordered feature names (critical for index mapping)
+        self.ordered_features = list(self.schema.feature_names)
+        
+        # 3. Names of categorical features
+        self.categorical_features = list(self.schema.categorical_feature_names)
+
+        _LOGGER.info(f"MiceImputer initialized. Found {len(self.cat_info)} categorical features to discretize.")
+        
+    def _post_process(self, imputed_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Applies schema-based discretization to a completed dataframe.
+        
+        This method works around the behavior of `discretize_categorical_values`
+        (which returns a full int32 array) by:
+        1. Calling it on the full, ordered feature array.
+        2. Extracting *only* the valid discretized categorical columns.
+        3. Updating the original float dataframe with these integer values.
+        """
+        # If no categorical features are defined, return the df as-is.
+        if not self.cat_info:
+            return imputed_df
+
+        try:
+            # 1. Ensure DataFrame columns match the schema order
+            # This is critical for the index-based categorical_info
+            df_ordered: pd.DataFrame = imputed_df[self.ordered_features] # type: ignore
+            
+            # 2. Convert to NumPy array
+            array_ordered = df_ordered.to_numpy()
+
+            # 3. Apply discretization utility (which returns a full int32 array)
+            # This array has *correct* categorical values but *truncated* continuous values.
+            discretized_array_int32 = discretize_categorical_values(
+                array_ordered,
+                self.cat_info,
+                start_at_zero=True  # Assuming 0-based indexing
+            )
+
+            # 4. Create a new DF from the int32 array, keeping the categorical columns.
+            df_discretized_cats = pd.DataFrame(
+                discretized_array_int32,
+                columns=self.ordered_features,
+                index=df_ordered.index  # <-- Critical: align index
+            )[self.categorical_features] # <-- Select only cat features
+
+            # 5. "Rejoin": Start with a fresh copy of the *original* imputed DF (which has correct continuous floats).
+            final_df = df_ordered.copy()
+            
+            # 6. Use .update() to "paste" the integer categorical values
+            # over the old float categorical values. Continuous floats are unaffected.
+            final_df.update(df_discretized_cats)
+            
+            return final_df
+
+        except Exception as e:
+            _LOGGER.error(f"Failed during post-processing discretization:\n\tInput DF shape: {imputed_df.shape}\n\tSchema features: {len(self.ordered_features)}\n\tCategorical info keys: {list(self.cat_info.keys())}\n{e}")
+            raise
+        
+    def _run_mice(self, 
+                  df: pd.DataFrame, 
+                  df_name: str) -> tuple[mf.ImputationKernel, list[pd.DataFrame], list[str]]:
+        """
+        Runs the MICE kernel and applies schema-based post-processing.
+        
+        Parameters:
+            df (pd.DataFrame): The input dataframe *with NaNs*. Should only contain feature columns.
+            df_name (str): The base name for the dataset.
+
+        Returns:
+            tuple[mf.ImputationKernel, list[pd.DataFrame], list[str]]:
+                - The trained MICE kernel
+                - A list of imputed and processed DataFrames
+                - A list of names for the new DataFrames
+        """
+        # Ensure input df only contains features from the schema and is in the correct order.
+        try:
+            df_feats = df[self.ordered_features]
+        except KeyError as e:
+            _LOGGER.error(f"Input DataFrame is missing required schema columns: {e}")
+            raise
+        
+        # 1. Initialize kernel
+        kernel = mf.ImputationKernel(
+            data=df_feats,
+            num_datasets=self.resulting_datasets,
+            random_state=self.random_state
+        )
+        
+        _LOGGER.info("➡️ Schema-based MICE imputation running...")
+        
+        # 2. Perform MICE
+        kernel.mice(self.iterations)
+        
+        # 3. Retrieve, process, and collect datasets
+        imputed_datasets = []
+        for i in range(self.resulting_datasets):
+            # complete_data returns a pd.DataFrame
+            completed_df = kernel.complete_data(dataset=i)
+            
+            # Apply our new discretization and ordering
+            processed_df = self._post_process(completed_df)
+            imputed_datasets.append(processed_df)
+
+        if not imputed_datasets:
+            _LOGGER.error("No imputed datasets were generated.")
+            raise ValueError()
+
+        # 4. Generate names
+        if self.resulting_datasets == 1:
+            imputed_dataset_names = [f"{df_name}_MICE"]
+        else:
+            imputed_dataset_names = [f"{df_name}_MICE_{i+1}" for i in range(self.resulting_datasets)]
+        
+        # 5. Validate indexes 
+        for imputed_df, subname in zip(imputed_datasets, imputed_dataset_names):
+            assert imputed_df.shape[0] == df.shape[0], f"❌ Row count mismatch in dataset {subname}"
+            assert all(imputed_df.index == df.index), f"❌ Index mismatch in dataset {subname}"
+        
+        _LOGGER.info("Schema-based MICE imputation complete.")
+        
+        return kernel, imputed_datasets, imputed_dataset_names
+        
+    def run_pipeline(self, 
+                     df_path_or_dir: Union[str,Path],
+                     save_datasets_dir: Union[str,Path], 
+                     save_metrics_dir: Union[str,Path],
+                     ):
+        """
+        Runs the complete MICE imputation pipeline.
+
+        This method automates the entire workflow:
+            1. Loads data from a CSV file path or a directory with CSV files.
+            2. Separates features and targets based on the `FeatureSchema`.
+            3. Runs the MICE algorithm on the feature set.
+            4. Applies schema-based post-processing to discretize categorical features.
+            5. Saves the final, processed, and imputed dataset(s) (re-joined with targets) to `save_datasets_dir`.
+            6. Generates and saves convergence and distribution plots for all imputed columns to `save_metrics_dir`.
+
+        Parameters
+        ----------
+        df_path_or_dir :[str,Path]
+            Path to a single CSV file or a directory containing multiple CSV files to impute.
+        save_datasets_dir : [str,Path]
+            Directory where the final imputed and processed dataset(s) will be saved as CSVs.
+        save_metrics_dir : [str,Path]
+            Directory where convergence and distribution plots will be saved.
+        """
+        # Check paths
+        save_datasets_path = make_fullpath(save_datasets_dir, make=True)
+        save_metrics_path = make_fullpath(save_metrics_dir, make=True)
+        
+        input_path = make_fullpath(df_path_or_dir)
+        if input_path.is_file():
+            all_file_paths = [input_path]
+        else:
+            all_file_paths = list(list_csv_paths(input_path).values())
+        
+        for df_path in all_file_paths:
+            
+            df, df_name = load_dataframe(df_path=df_path, kind="pandas")
+            
+            df_features: pd.DataFrame = df[self.schema.feature_names] # type: ignore
+            df_targets = df.drop(columns=self.schema.feature_names)
+            
+            imputed_column_names = _get_na_column_names(df=df_features)
+            
+            kernel, imputed_datasets, imputed_dataset_names = self._run_mice(df=df_features, df_name=df_name)
+            
+            save_imputed_datasets(save_dir=save_datasets_path, imputed_datasets=imputed_datasets, df_targets=df_targets, imputed_dataset_names=imputed_dataset_names)
+            
+            get_convergence_diagnostic(kernel=kernel, imputed_dataset_names=imputed_dataset_names, column_names=imputed_column_names, root_dir=save_metrics_path)
+            
+            get_imputed_distributions(kernel=kernel, df_name=df_name, root_dir=save_metrics_path, column_names=imputed_column_names)
+
+
 def info():
     _script_info(__all__)

ml_tools/ML_utilities.py

@@ -1,18 +1,24 @@
 import pandas as pd
 from pathlib import Path
-from typing import Union, Any, Optional
+from typing import Union, Any, Optional, Dict, List, Iterable
+import torch
+from torch import nn
+
 
 from .path_manager import make_fullpath, list_subdirectories, list_files_by_extension
 from ._script_info import _script_info
 from ._logger import _LOGGER
-from .keys import DatasetKeys, PytorchModelArchitectureKeys, PytorchArtifactPathKeys, SHAPKeys
+from .keys import DatasetKeys, PytorchModelArchitectureKeys, PytorchArtifactPathKeys, SHAPKeys, UtilityKeys, PyTorchCheckpointKeys
 from .utilities import load_dataframe
-from .custom_logger import save_list_strings
+from .custom_logger import save_list_strings, custom_logger
 
 
 __all__ = [
     "find_model_artifacts",
-    "select_features_by_shap"
+    "select_features_by_shap",
+    "get_model_parameters",
+    "inspect_pth_file",
+    "set_parameter_requires_grad"
 ]
 
 
@@ -226,5 +232,248 @@ def select_features_by_shap(
     return final_features
 
 
+def get_model_parameters(model: nn.Module, save_dir: Optional[Union[str,Path]]=None) -> Dict[str, int]:
+    """
+    Calculates the total and trainable parameters of a PyTorch model.
+
+    Args:
+        model (nn.Module): The PyTorch model to inspect.
+        save_dir: Optional directory to save the output as a JSON file.
+
+    Returns:
+        Dict[str, int]: A dictionary containing:
+            - "total_params": The total number of parameters.
+            - "trainable_params": The number of trainable parameters (where requires_grad=True).
+    """
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+    report = {
+        UtilityKeys.TOTAL_PARAMS: total_params,
+        UtilityKeys.TRAINABLE_PARAMS: trainable_params
+    }
+    
+    if save_dir is not None:
+        output_dir = make_fullpath(save_dir, make=True, enforce="directory")
+        custom_logger(data=report,
+                      save_directory=output_dir,
+                      log_name=UtilityKeys.MODEL_PARAMS_FILE,
+                      dict_as="json")
+
+    return report
+
+
+def inspect_pth_file(
+    pth_path: Union[str, Path],
+    save_dir: Union[str, Path],
+) -> None:
+    """
+    Inspects a .pth file (e.g., checkpoint) and saves a human-readable
+    JSON summary of its contents.
+
+    Args:
+        pth_path (str | Path): The path to the .pth file to inspect.
+        save_dir (str | Path): The directory to save the JSON report.
+
+    Returns:
+        Dict (str, Any): A dictionary containing the inspection report.
+
+    Raises:
+        ValueError: If the .pth file is empty or in an unrecognized format.
+    """
+    # --- 1. Validate paths ---
+    pth_file = make_fullpath(pth_path, enforce="file")
+    output_dir = make_fullpath(save_dir, make=True, enforce="directory")
+    pth_name = pth_file.stem
+
+    # --- 2. Load data ---
+    try:
+        # Load onto CPU to avoid GPU memory issues
+        loaded_data = torch.load(pth_file, map_location=torch.device('cpu'))
+    except Exception as e:
+        _LOGGER.error(f"Failed to load .pth file '{pth_file}': {e}")
+        raise
+
+    # --- 3. Initialize Report ---
+    report = {
+        "top_level_type": str(type(loaded_data)),
+        "top_level_summary": {},
+        "model_state_analysis": None,
+        "notes": []
+    }
+
+    # --- 4. Parse loaded data ---
+    if isinstance(loaded_data, dict):
+        # --- Case 1: Loaded data is a dictionary (most common case) ---
+        # "main loop" that iterates over *everything* first.
+        for key, value in loaded_data.items():
+            key_summary = {}
+            val_type = str(type(value))
+            key_summary["type"] = val_type
+            
+            if isinstance(value, torch.Tensor):
+                key_summary["shape"] = list(value.shape)
+                key_summary["dtype"] = str(value.dtype)
+            elif isinstance(value, dict):
+                key_summary["key_count"] = len(value)
+                key_summary["key_preview"] = list(value.keys())[:5]
+            elif isinstance(value, (int, float, str, bool)):
+                key_summary["value_preview"] = str(value)
+            elif isinstance(value, (list, tuple)):
+                 key_summary["value_preview"] = str(value)[:100]
+            
+            report["top_level_summary"][key] = key_summary
+
+        # Now, try to find the model state_dict within the dict
+        if PyTorchCheckpointKeys.MODEL_STATE in loaded_data and isinstance(loaded_data[PyTorchCheckpointKeys.MODEL_STATE], dict):
+            report["notes"].append(f"Found standard checkpoint key: '{PyTorchCheckpointKeys.MODEL_STATE}'. Analyzing as model state_dict.")
+            state_dict = loaded_data[PyTorchCheckpointKeys.MODEL_STATE]
+            report["model_state_analysis"] = _generate_weight_report(state_dict)
+        
+        elif all(isinstance(v, torch.Tensor) for v in loaded_data.values()):
+            report["notes"].append("File dictionary contains only tensors. Analyzing entire dictionary as model state_dict.")
+            state_dict = loaded_data
+            report["model_state_analysis"] = _generate_weight_report(state_dict)
+        
+        else:
+            report["notes"].append("Could not identify a single model state_dict. See top_level_summary for all contents. No detailed weight analysis will be performed.")
+
+    elif isinstance(loaded_data, nn.Module):
+        # --- Case 2: Loaded data is a full pickled model ---
+        # _LOGGER.warning("Loading a full, pickled nn.Module is not recommended. Inspecting its state_dict().")
+        report["notes"].append("File is a full, pickled nn.Module. This is not recommended. Extracting state_dict() for analysis.")
+        state_dict = loaded_data.state_dict()
+        report["model_state_analysis"] = _generate_weight_report(state_dict)
+
+    else:
+        # --- Case 3: Unrecognized format (e.g., single tensor, list) ---
+        _LOGGER.error(f"Could not parse .pth file. Loaded data is of type {type(loaded_data)}, not a dict or nn.Module.")
+        raise ValueError()
+
+    # --- 5. Save Report ---
+    custom_logger(data=report,
+                  save_directory=output_dir,
+                  log_name=UtilityKeys.PTH_FILE + pth_name,
+                  dict_as="json")
+
+
+def _generate_weight_report(state_dict: dict) -> dict:
+    """
+    Internal helper to analyze a state_dict and return a structured report.
+    
+    Args:
+        state_dict (dict): The model state_dict to analyze.
+
+    Returns:
+        dict: A report containing total parameters and a per-parameter breakdown.
+    """
+    weight_report = {}
+    total_params = 0
+    if not isinstance(state_dict, dict):
+        _LOGGER.warning(f"Attempted to generate weight report on non-dict type: {type(state_dict)}")
+        return {"error": "Input was not a dictionary."}
+
+    for key, tensor in state_dict.items():
+        if not isinstance(tensor, torch.Tensor):
+             _LOGGER.warning(f"Skipping key '{key}' in state_dict: value is not a tensor (type: {type(tensor)}).")
+             weight_report[key] = {
+                 "type": str(type(tensor)),
+                 "value_preview": str(tensor)[:50] # Show a preview
+             }
+             continue
+        weight_report[key] = {
+            "shape": list(tensor.shape),
+            "dtype": str(tensor.dtype),
+            "requires_grad": tensor.requires_grad,
+            "num_elements": tensor.numel()
+        }
+        total_params += tensor.numel()
+
+    return {
+        "total_parameters": total_params,
+        "parameter_key_count": len(weight_report),
+        "parameters": weight_report
+    }
+
+
+def set_parameter_requires_grad(
+    model: nn.Module,
+    unfreeze_last_n_params: int,
+) -> int:
+    """
+    Freezes or unfreezes parameters in a model based on unfreeze_last_n_params.
+
+    - N = 0: Freezes ALL parameters.
+    - N > 0 and N < total: Freezes ALL parameters, then unfreezes the last N.
+    - N >= total: Unfreezes ALL parameters.
+
+    Note: 'N' refers to individual parameter tensors (e.g., `layer.weight`
+    or `layer.bias`), not modules or layers. For example, to unfreeze
+    the final nn.Linear layer, you would use N=2 (for its weight and bias).
+
+    Args:
+        model (nn.Module): The model to modify.
+        unfreeze_last_n_params (int):
+            The number of parameter tensors to unfreeze, starting from
+            the end of the model.
+
+    Returns:
+        int: The total number of individual parameters (elements) that were set to `requires_grad=True`.
+    """
+    if unfreeze_last_n_params < 0:
+        _LOGGER.error(f"unfreeze_last_n_params must be >= 0, but got {unfreeze_last_n_params}")
+        raise ValueError()
+
+    # --- Step 1: Get all parameter tensors ---
+    all_params = list(model.parameters())
+    total_param_tensors = len(all_params)
+
+    # --- Case 1: N = 0 (Freeze ALL parameters) ---
+    # early exit for the "freeze all" case.
+    if unfreeze_last_n_params == 0:
+        params_frozen = _set_params_grad(all_params, requires_grad=False)
+        _LOGGER.warning(f"Froze all {total_param_tensors} parameter tensors ({params_frozen} total elements).")
+        return 0  # 0 parameters unfrozen
+
+    # --- Case 2: N >= total (Unfreeze ALL parameters) ---
+    if unfreeze_last_n_params >= total_param_tensors:
+        if unfreeze_last_n_params > total_param_tensors:
+            _LOGGER.warning(f"Requested to unfreeze {unfreeze_last_n_params} params, but model only has {total_param_tensors}. Unfreezing all.")
+        
+        params_unfrozen = _set_params_grad(all_params, requires_grad=True)
+        _LOGGER.info(f"Unfroze all {total_param_tensors} parameter tensors ({params_unfrozen} total elements) for training.")
+        return params_unfrozen
+
+    # --- Case 3: 0 < N < total (Standard: Freeze all, unfreeze last N) ---
+    # Freeze ALL
+    params_frozen = _set_params_grad(all_params, requires_grad=False)
+    _LOGGER.info(f"Froze {params_frozen} parameters.")
+
+    # Unfreeze the last N
+    params_to_unfreeze = all_params[-unfreeze_last_n_params:]
+    
+    # these are all False, so the helper will set them to True
+    params_unfrozen = _set_params_grad(params_to_unfreeze, requires_grad=True)
+
+    _LOGGER.info(f"Unfroze the last {unfreeze_last_n_params} parameter tensors ({params_unfrozen} total elements) for training.")
+
+    return params_unfrozen
+
+
+def _set_params_grad(
+    params: Iterable[nn.Parameter], 
+    requires_grad: bool
+) -> int:
+    """
+    A helper function to set the `requires_grad` attribute for an iterable
+    of parameters and return the total number of elements changed.
+    """
+    params_changed = 0
+    for param in params:
+        if param.requires_grad != requires_grad:
+            param.requires_grad = requires_grad
+            params_changed += param.numel()
+    return params_changed
+
 def info():
     _script_info(__all__)

ml_tools/custom_logger.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 from datetime import datetime
-from typing import Union, List, Dict, Any
+from typing import Union, List, Dict, Any, Literal
 import traceback
 import json
 import csv
@@ -29,6 +29,7 @@ def custom_logger(
     ],
     save_directory: Union[str, Path],
     log_name: str,
+    dict_as: Literal['auto', 'json', 'csv'] = 'auto',
 ) -> None:
     """
     Logs various data types to corresponding output formats:
@@ -36,10 +37,10 @@ def custom_logger(
     - list[Any]                    → .txt
         Each element is written on a new line.
 
-    - dict[str, list[Any]]        → .csv
+    - dict[str, list[Any]]        → .csv    (if dict_as='auto' or 'csv')
         Dictionary is treated as tabular data; keys become columns, values become rows.
 
-    - dict[str, scalar]           → .json
+    - dict[str, scalar]           → .json   (if dict_as='auto' or 'json')
         Dictionary is treated as structured data and serialized as JSON.
 
     - str                         → .log
@@ -52,26 +53,43 @@ def custom_logger(
         data: The data to be logged. Must be one of the supported types.
         save_directory: Directory where the log will be saved. Created if it does not exist.
         log_name: Base name for the log file. Timestamp will be appended automatically.
+        dict_as ('auto'|'json'|'csv'): 
+            - 'auto': Guesses format (JSON or CSV) based on dictionary content.
+            - 'json': Forces .json format for any dictionary.
+            - 'csv': Forces .csv format. Will fail if dict values are not all lists.
 
     Raises:
         ValueError: If the data type is unsupported.
     """
     try:
+        if not isinstance(data, BaseException) and not data:
+            _LOGGER.warning("Empty data received. No log file will be saved.")
+            return
+        
         save_path = make_fullpath(save_directory, make=True)
         
         timestamp = datetime.now().strftime(r"%Y%m%d_%H%M%S")
         log_name = sanitize_filename(log_name)
         
         base_path = save_path / f"{log_name}_{timestamp}"
-
+        
+        # Router
         if isinstance(data, list):
             _log_list_to_txt(data, base_path.with_suffix(".txt"))
 
         elif isinstance(data, dict):
-            if all(isinstance(v, list) for v in data.values()):
-                _log_dict_to_csv(data, base_path.with_suffix(".csv"))
-            else:
+            if dict_as == 'json':
                 _log_dict_to_json(data, base_path.with_suffix(".json"))
+            
+            elif dict_as == 'csv':
+                # This will raise a ValueError if data is not all lists
+                _log_dict_to_csv(data, base_path.with_suffix(".csv"))
+            
+            else: # 'auto' mode
+                if all(isinstance(v, list) for v in data.values()):
+                    _log_dict_to_csv(data, base_path.with_suffix(".csv"))
+                else:
+                    _log_dict_to_json(data, base_path.with_suffix(".json"))
 
         elif isinstance(data, str):
             _log_string_to_log(data, base_path.with_suffix(".log"))
@@ -83,7 +101,7 @@ def custom_logger(
             _LOGGER.error("Unsupported data type. Must be list, dict, str, or BaseException.")
             raise ValueError()
 
-        _LOGGER.info(f"Log saved to: '{base_path}'")
+        _LOGGER.info(f"Log saved as: '{base_path.name}'")
 
     except Exception:
         _LOGGER.exception(f"Log not saved.")

ml_tools/keys.py

@@ -80,6 +80,14 @@ class PyTorchCheckpointKeys:
     BEST_SCORE = "best_score"
 
 
+class UtilityKeys:
+    """Keys used for utility modules"""
+    MODEL_PARAMS_FILE = "model_parameters"
+    TOTAL_PARAMS = "Total Parameters"
+    TRAINABLE_PARAMS = "Trainable Parameters"
+    PTH_FILE = "pth report "
+
+
 class _OneHotOtherPlaceholder:
     """Used internally by GUI_tools."""
     OTHER_GUI = "OTHER"

ml_tools/utilities.py

@@ -7,16 +7,19 @@ from typing import Literal, Union, Optional, Any, Iterator, Tuple, overload
 from .path_manager import sanitize_filename, make_fullpath, list_csv_paths
 from ._script_info import _script_info
 from ._logger import _LOGGER
+from ._schema import FeatureSchema
 
 
 # Keep track of available tools
 __all__ = [
     "load_dataframe",
     "load_dataframe_greedy",
+    "load_dataframe_with_schema",
     "yield_dataframes_from_dir",
     "merge_dataframes",
     "save_dataframe_filename",
     "save_dataframe",
+    "save_dataframe_with_schema",
     "distribute_dataset_by_target",
     "train_dataset_orchestrator",
     "train_dataset_yielder"
@@ -174,6 +177,68 @@ def load_dataframe_greedy(directory: Union[str, Path],
     return df
 
 
+def load_dataframe_with_schema(
+    df_path: Union[str, Path], 
+    schema: "FeatureSchema",
+    all_strings: bool = False,
+) -> Tuple[pd.DataFrame, str]:
+    """
+    Loads a CSV file into a Pandas DataFrame, strictly validating its
+    feature columns against a FeatureSchema.
+
+    This function wraps `load_dataframe`. After loading, it validates
+    that the first N columns of the DataFrame (where N =
+    len(schema.feature_names)) contain *exactly* the set of features
+    specified in the schema.
+
+    - If the columns are present but out of order, they are reordered.
+    - If any required feature is missing from the first N columns, it fails.
+    - If any extra column is found within the first N columns, it fails.
+
+    Columns *after* the first N are considered target columns and are
+    logged for verification.
+
+    Args:
+        df_path (str, Path): 
+            The path to the CSV file.
+        schema (FeatureSchema): 
+            The schema object to validate against.
+        all_strings (bool): 
+            If True, loads all columns as string data types.
+
+    Returns:
+        (Tuple[pd.DataFrame, str]):
+            A tuple containing the loaded, validated (and possibly
+            reordered) pandas DataFrame and the base name of the file.
+            
+    Raises:
+        ValueError: 
+            - If the DataFrame is missing columns required by the schema
+              within its first N columns.
+            - If the DataFrame's first N columns contain unexpected
+              columns that are not in the schema.
+        FileNotFoundError: 
+            If the file does not exist at the given path.
+    """
+    # Step 1: Load the dataframe using the original function
+    try:
+        df, df_name = load_dataframe(
+            df_path=df_path, 
+            use_columns=None,  # Load all columns for validation
+            kind="pandas", 
+            all_strings=all_strings,
+            verbose=True
+        )
+    except Exception as e:
+        _LOGGER.error(f"Failed during initial load for schema validation: {e}")
+        raise e
+    
+    # Step 2: Call the helper to validate and reorder
+    df_validated = _validate_and_reorder_schema(df=df, schema=schema)
+
+    return df_validated, df_name
+
+
 def yield_dataframes_from_dir(datasets_dir: Union[str,Path], verbose: bool=True):
     """
     Iterates over all CSV files in a given directory, loading each into a Pandas DataFrame.
@@ -330,6 +395,52 @@ def save_dataframe(df: Union[pd.DataFrame, pl.DataFrame], full_path: Path):
                             filename=full_path.name)
 
 
+def save_dataframe_with_schema(
+    df: pd.DataFrame, 
+    full_path: Path,
+    schema: "FeatureSchema"
+) -> None:
+    """
+    Saves a pandas DataFrame to a CSV, strictly enforcing that the
+    first N columns match the FeatureSchema.
+
+    This function validates that the first N columns of the DataFrame
+    (where N = len(schema.feature_names)) contain *exactly* the set
+    of features specified in the schema.
+    
+    - If the columns are present but out of order, they are reordered.
+    - If any required feature is missing from the first N columns, it fails.
+    - If any extra column is found within the first N columns, it fails.
+
+    Columns *after* the first N are considered target columns and are
+    logged for verification.
+
+    Args:
+        df (pd.DataFrame): 
+            The DataFrame to save.
+        full_path (Path): 
+            The complete file path where the DataFrame will be saved.
+        schema (FeatureSchema): 
+            The schema object to validate against.
+
+    Raises:
+        ValueError: 
+            - If the DataFrame is missing columns required by the schema
+              within its first N columns.
+            - If the DataFrame's first N columns contain unexpected
+              columns that are not in the schema.
+    """
+    if not isinstance(full_path, Path) or not full_path.suffix.endswith(".csv"):
+        _LOGGER.error('A path object pointing to a .csv file must be provided.')
+        raise ValueError()
+    
+    # Call the helper to validate and reorder
+    df_to_save = _validate_and_reorder_schema(df=df, schema=schema)
+    
+    # Call the original save function
+    save_dataframe(df=df_to_save, full_path=full_path)
+
+
 def distribute_dataset_by_target(
     df_or_path: Union[pd.DataFrame, str, Path],
     target_columns: list[str],
@@ -442,5 +553,72 @@ def train_dataset_yielder(
         yield (df_features, df_target, feature_names, target_col)
 
 
+def _validate_and_reorder_schema(
+    df: pd.DataFrame, 
+    schema: "FeatureSchema"
+) -> pd.DataFrame:
+    """
+    Internal helper to validate and reorder a DataFrame against a schema.
+
+    Checks for missing, extra, and out-of-order feature columns
+    (the first N columns). Returns a reordered DataFrame if necessary.
+    Logs all actions.
+
+    Raises:
+        ValueError: If validation fails.
+    """
+    # Get schema and DataFrame column info
+    expected_features = list(schema.feature_names)
+    expected_set = set(expected_features)
+    n_features = len(expected_features)
+    
+    all_df_columns = df.columns.to_list()
+
+    # --- Strict Validation ---
+
+    # 0. Check if DataFrame is long enough
+    if len(all_df_columns) < n_features:
+        _LOGGER.error(f"DataFrame has only {len(all_df_columns)} columns, but schema requires {n_features} features.")
+        raise ValueError()
+    
+    df_feature_cols = all_df_columns[:n_features]
+    df_feature_set = set(df_feature_cols)
+    df_target_cols = all_df_columns[n_features:]
+
+    # 1. Check for missing features
+    missing_from_df = expected_set - df_feature_set
+    if missing_from_df:
+        _LOGGER.error(f"DataFrame's first {n_features} columns are missing required schema features: {missing_from_df}")
+        raise ValueError()
+
+    # 2. Check for extra (unexpected) features
+    extra_in_df = df_feature_set - expected_set
+    if extra_in_df:
+        _LOGGER.error(f"DataFrame's first {n_features} columns contain unexpected columns: {extra_in_df}")
+        raise ValueError()
+
+    # --- Reordering ---
+    
+    df_to_process = df
+
+    # If we pass validation, the sets are equal. Now check order.
+    if df_feature_cols == expected_features:
+        _LOGGER.info("DataFrame feature columns already match schema order.")
+    else:
+        _LOGGER.warning("DataFrame feature columns do not match schema order. Reordering...")
+        
+        # Rebuild the DataFrame with the correct feature order + target columns
+        new_order = expected_features + df_target_cols
+        df_to_process = df[new_order]
+
+    # Log the presumed target columns for user verification
+    if not df_target_cols:
+        _LOGGER.warning(f"No target columns were found after index {n_features-1}.")
+    else:
+        _LOGGER.info(f"Presumed Target Columns: {df_target_cols}")
+    
+    return df_to_process # type: ignore
+
+
 def info():
     _script_info(__all__)