Model/model_data/simulation_system/training.py

import argparse
# import boto3
import os 
from pathlib import Path
from datetime import datetime
from typing import List
from core.Logger import logger
from core.DataLoader import DataLoader
from core.FeatureProcessor import FeatureProcessor
from MLModel.Models import AutogluonModel
import pandas as pd
from core.Settings import (
    MODEL_DIRECTORY,
    BASE_REGISTRY_PATH,
    REGISTRY_FILE,
    MODEL_FOLDER,
    METRICS_FOLDER,
    DEPLOYMENT_FOLDER,
    SUBSAMPLE_FACTOR,
    MODEL_HYPERPARAMETERS
)
import seaborn as sns
import matplotlib.pyplot as plt

TIMESTAMP = datetime.now().strftime(format="%Y-%m-%d_%H-%M-%S")     

# FOR TESTING
# train_filepath = "./model_build_data/change_data/rdsap_full/train_validation_data.parquet"
# test_filepath = "./model_build_data/change_data/rdsap_full/test_data.parquet"
# target_column = "RDSAP_CHANGE"
# model_type = "autogluon"
# hyperparameter = HYPERPARAMETERS
# SUBSAMPLE_FACTOR = 200

# SESSION = boto3.Session()

# S3_CLIENT = SESSION.client(
#     service_name="s3",
#     aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID", 'admin'),
#     aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY", 'password'),
#     endpoint_url=os.environ.get("ENDPOINT_URL", "http://localhost:9000")
# )

# S3_CLIENT.create_bucket
# S3_CLIENT.list_buckets()

def ingest_arguments() -> argparse.Namespace:
    """
    Helper function to take in arguments from script start
    """

    parser = argparse.ArgumentParser(description='Inputs for training script')

    parser.add_argument('--train-filepath', type=str, help='Location of Parquet dataset to load for training', required=True)
    parser.add_argument('--test-filepath', type=str, help='Location of Parquet dataset to load for testing', required=True)
    parser.add_argument('--model-type', type=str, help='The type of model to train', choices=["autogluon"], default="autogluon")
    parser.add_argument('--target-column', type=str, help='The response variable', choices=["RDSAP_CHANGE", "HEAT_DEMAND_CHANGE"], default='RDSAP_CHANGE')

    args = parser.parse_args()

    return args
            

def training(
        train_filepath: str, 
        test_filepath: str, 
        target_column: str = "RDSAP_CHANGE", 
        model_type: str = "autogluon", 
        hyperparameters: dict = None
        ) -> None:
    """
    Pipeline to run training on the dataset
    """

    logger.info('--- Loading data ---')
    dataloader = DataLoader()
    train_df = dataloader.load(filepath=train_filepath)
    test_df = dataloader.load(filepath=test_filepath)
 
    logger.info('--- Feature processing ---')

    feature_processor = FeatureProcessor()

    subsample_amount = round(len(train_df)/SUBSAMPLE_FACTOR)

    train_df = feature_processor.process(train_df, target_column=target_column, subsample_amount=subsample_amount)
    test_df = feature_processor.process(test_df, target_column=target_column)

    logger.info('--- Build Model ---')

    logger.info("--- Load Hyperparameters ---")

    if hyperparameters is None:
        logger.info("Use base hyperparameters in settings")
        hyperparameters = MODEL_HYPERPARAMETERS[model_type]
        logger.info(f'Hyperparameters are: {hyperparameters}')

    if model_type == "autogluon":
        model_root = f"{target_column}-{hyperparameters['presets']}-{hyperparameters['time_limit']}-{TIMESTAMP}".lower()
        output_base = Path(MODEL_DIRECTORY) / target_column / model_type / model_root 

        model = AutogluonModel(
            output_filepath = output_base / MODEL_FOLDER
            )
    else:
        logger.error("No alternative model implemented yet")
        exit(1)
    
    model.train_model(
        data=train_df, 
        target_column=target_column, 
        hyperparameters=hyperparameters
        )
    
    logger.info("--- Save Model ---")
    model.save_model(output_filepath=model.output_filepath)

    logger.info('--- Generate evaluation metrics ---')
    metrics_df = model.model_evaluation(
        validation_data=test_df, 
        target_column=target_column,
        metrics_location = output_base / METRICS_FOLDER
        )
    
    logger.info("--- Generate metric outputs using predictions ---")
    # TODO: can have a model.metric_outputs method
    # FOr not just do it here
    residual_df = pd.DataFrame(list(zip(test_df[target_column], model.predictions)), columns=['true', 'pred'])
    
    # image formatting
    # TODO: move to settings file , AXIS_FONT, TITLE_FONT
    axis_fs = 18 #fontsize
    title_fs = 22 #fontsize
    sns.set(style="whitegrid")
    ax = sns.scatterplot(x="true", y="pred",data=residual_df)
    ax.set_aspect('equal')
    ax.set_xlabel(f'True {target_column}',fontsize = axis_fs) 
    ax.set_ylabel(f'Predicted {target_column}', fontsize = axis_fs)#ylabel
    ax.set_title('Residuals', fontsize = title_fs)

    # Square aspect ratio
    ax.plot([-100, 100], [-100, 100], 'black', linewidth=1)

    plt.tight_layout()
    RESIDUAL_FILE = "residuals.png"
    plt.savefig(output_base / METRICS_FOLDER / RESIDUAL_FILE, dpi=120) 

    # TODO: for cml, we might want to have class that outputs all data and plots to add to the report
    # If we want residual plot/ any plots, we will need to self host
    # plt.savefig(RESIDUAL_FILE, dpi=120) 

    
    # TODO: introduce a seperate script for model optimisation, and from there, optimise for deployment
    # Imagining for now that the model trained here is the best model amongst all models built

    logger.info("--- Optimising model for deployment ---")

    deployment_model_path = model.optimise_model_for_deployment(deployment_path= output_base / DEPLOYMENT_FOLDER)
    logger.info(f"Optimised version of best model can be found at: {deployment_model_path}")

    # TODO: Need a model registry - for now have this as a CSV
    # Save this in the model directory
    logger.info("--- Append registry with new model ---")

    registry_path = BASE_REGISTRY_PATH / target_column / REGISTRY_FILE

    if registry_path.exists():
        logger.info("Registry file found - Loading into Dataframe")
        registry_df = pd.read_csv(registry_path, index_col=None)
    else:
        # TODO: Moved columns into settings: MODEL_DETAILS and Metrics class columns 
        registry_df = pd.DataFrame(columns=['model_type', 'model_name', 'model_location', 'mean_absolute_error', 'root_mean_squared_error', 'mean_squared_error', 'r2', 'pearsonr', 'median_absolute_error', 'mape', 'best_model'])

    model_details_df = pd.DataFrame(
        [{
            'model_type': model_type, 
            'model_name': model_root, 
            'model_location': deployment_model_path
        }]
        )
    
    registry_row = pd.concat([model_details_df, metrics_df], axis=1)
    registry_df = pd.concat([registry_df, registry_row], axis=0).reset_index(drop=True)

    # TODO: will need a rebuild script metric script -i.e. if we add new metrics, we will want to load models and regenerate new metrics
    # TODO: decide metric to optimise to
    registry_df = registry_df.sort_values("mean_absolute_error", ascending=False).reset_index(drop=True)
    registry_df['best_model'] = [False]*len(registry_df)
    registry_df.loc[0, 'best_model'] = True

    logger.info("--- Saving new model to registry ---")
    registry_df.to_csv(registry_path, index=False)

    logger.info("--- Training Pipeline Complete --- ")


if __name__ == "__main__":

    logger.info('---Begin Pipeline---')

    logger.info('---Ingest Arguments---')
    args = ingest_arguments()

    # To run script: python3 training.py --train-filepath ./model_build_data/change_data/rdsap_full/train_validation_data.parquet --test-filepath ./model_build_data/change_data/rdsap_full/test_data.parquet
    # TODO: Ingest hyper parameters from somewhere - currently change at the top of script
    training(
        train_filepath=args.train_filepath, 
        test_filepath=args.test_filepath, 
        target_column=args.target_column, 
        model_type=args.model_type
        )