# Import libraries
import argparse
from azureml.core import Run
import joblib
import json
import os
import pandas as pd
import shutil

# Import functions from train.py
from train import split_data, train_model, get_model_metrics

# Get the output folder for the model from the '--output_folder' parameter
parser = argparse.ArgumentParser()
parser.add_argument('--output_folder', type=str, dest='output_folder', default="outputs")

args = parser.parse_args()
print(args)

output_folder = args.output_folder

# Get the experiment run context
run = Run.get_context()

# load the safe driver prediction dataset
train_df = pd.read_csv('porto_seguro_safe_driver_prediction_input.csv')

# Load the parameters for training the model from the file
with open("parameters.json") as f:
    pars = json.load(f)
    parameters = pars["training"]

# Log each of the parameters to the run
for param_name, param_value in parameters.items():
    run.log(param_name, param_value)

# Call the functions defined in this file
train_data, valid_data = split_data(train_df)
data = [train_data, valid_data]
model = train_model(data, parameters)

# Print the resulting metrics for the model
model_metrics = get_model_metrics(model, data)
print(model_metrics)

for k, v in model_metrics.items():
    run.log(k, v)

# Save the trained model to the output folder
os.makedirs(output_folder, exist_ok=True)
output_path = output_folder + "/porto_seguro_safe_driver_model.pkl"
joblib.dump(value=model, filename=output_path)

run.complete()

[((14, 9, 14, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((23, 6, 23, 23), 'azureml.core.Run.get_context', 'Run.get_context', ({}, {}), '()', False, 'from azureml.core import Run\n'), ((26, 11, 26, 71), 'pandas.read_csv', 'pd.read_csv', ({(26, 23, 26, 70): '"""porto_seguro_safe_driver_prediction_input.csv"""'}, {}), "('porto_seguro_safe_driver_prediction_input.csv')", True, 'import pandas as pd\n'), ((38, 25, 38, 45), 'train.split_data', 'split_data', ({(38, 36, 38, 44): 'train_df'}, {}), '(train_df)', False, 'from train import split_data, train_model, get_model_metrics\n'), ((40, 8, 40, 37), 'train.train_model', 'train_model', ({(40, 20, 40, 24): 'data', (40, 26, 40, 36): 'parameters'}, {}), '(data, parameters)', False, 'from train import split_data, train_model, get_model_metrics\n'), ((43, 16, 43, 46), 'train.get_model_metrics', 'get_model_metrics', ({(43, 34, 43, 39): 'model', (43, 41, 43, 45): 'data'}, {}), '(model, data)', False, 'from train import split_data, train_model, get_model_metrics\n'), ((50, 0, 50, 41), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((52, 0, 52, 46), 'joblib.dump', 'joblib.dump', (), '', False, 'import joblib\n'), ((30, 11, 30, 23), 'json.load', 'json.load', ({(30, 21, 30, 22): 'f'}, {}), '(f)', False, 'import json\n')]