Upload train_model_main.py

train_model_main.py

@@ -0,0 +1,254 @@
+import argparse
+import pandas as pd
+import tensorflow as tf
+import numpy as np
+import matplotlib.pyplot as plt
+
+# from numpy.random import seed
+import random
+import os
+
+import pickle
+import shap
+import dill
+
+from utils import encode_categorical, scale_numerical, NoPhysicsModels, unpickle_file
+from alloy_data_preprocessing import add_physics_features
+import tensorflow as tf
+from tensorflow.keras import initializers
+from sklearn.model_selection import train_test_split
+from sklearn.ensemble import RandomForestRegressor
+from sklearn.metrics import mean_absolute_error, mean_absolute_percentage_error, r2_score
+
+SEED = 42
+
+
+def set_all_seeds(seed=SEED):
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    tf.keras.utils.set_random_seed(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def setup_model(num_outputs):
+    model = tf.keras.models.Sequential(
+        [
+            tf.keras.layers.Flatten(),
+            tf.keras.layers.Dense(
+                8,
+                kernel_initializer=initializers.RandomNormal(stddev=0.00001),  # Initially was at 0.01
+                bias_initializer=initializers.Zeros(),
+                activation="relu",
+            ),
+            tf.keras.layers.Dense(
+                4,
+                activation="relu",
+                kernel_initializer=initializers.RandomNormal(stddev=0.00001),  # Initially was at 0.01
+                bias_initializer=initializers.Zeros(),
+            ),
+            tf.keras.layers.Dense(
+                num_outputs,
+                activation="relu",
+                kernel_initializer=initializers.RandomNormal(stddev=0.00001),  # Initially was at 0.01
+                bias_initializer=initializers.Zeros(),
+            ),
+        ]
+    )
+    return model
+
+
+def prepare_data(data, columns_num, columns_target, main_folder, data_type="path", seed=SEED):
+    # Create folder if doesn't exist
+    if not os.path.exists(main_folder):
+        os.makedirs(main_folder)
+
+    columns_numerical = columns_num.copy()
+
+    ### Read data
+    print(data_type)
+    if data_type == "path":
+        df = pd.read_csv(data, sep=";")
+    else:
+        df = data.copy()
+
+    ### Remove columns not used during training
+    X = df.drop(columns=columns_target)
+    y = df[columns_target]
+
+    # Remove the index columns (if coming from the sampling pipeline)
+    if "Index" in X.columns:
+        X.drop(columns=["Index"], inplace=True)
+
+    ### Get categorical columns
+    columns_categorical = [column for column in X.columns if column not in columns_numerical]
+    # ### Remove target from column names
+    # for target in columns_target:
+    #     columns_numerical.remove(target)
+    print("lllllllllllllllllllllllll")
+    print(X.columns)
+
+    ### Encode variables into one-hot
+    X, one_hot_scaler = encode_categorical(X, columns_categorical)
+    X, minmax_scaler_inputs = scale_numerical(
+        X, [column for column in columns_numerical if column not in columns_target]
+    )
+    y, minmax_scaler_targets = scale_numerical(y, columns_target)
+
+    ### Split data
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=seed)
+
+    ### Pickle data
+    with open(os.path.join(main_folder, f"X_test_data.pickle"), "wb+") as file:
+        pickle.dump(X_test, file)
+
+    with open(os.path.join(main_folder, f"y_test_data.pickle"), "wb+") as file:
+        pickle.dump(y_test, file)
+
+    with open(os.path.join(main_folder, f"one_hot_scaler.pickle"), "wb+") as file:
+        pickle.dump(one_hot_scaler, file)
+
+    with open(os.path.join(main_folder, f"minmax_scaler_inputs.pickle"), "wb+") as file:
+        pickle.dump(minmax_scaler_inputs, file)
+
+    with open(os.path.join(main_folder, f"minmax_scaler_targets.pickle"), "wb+") as file:
+        pickle.dump(minmax_scaler_targets, file)
+
+    return X_train, X_test, y_train, y_test
+
+
+def train_model_ml(X_train, X_test, y_train, y_test, main_folder, model_path, seed=SEED):
+    set_all_seeds(seed)
+
+    model = RandomForestRegressor(random_state=seed)
+    model.fit(X_train, y_train)
+    y_hat = model.predict(X_test)
+    print("----------------")
+    print("Model performance")
+    print("MAE", mean_absolute_error(y_test, y_hat))
+    print("MAPE", mean_absolute_percentage_error(y_test, y_hat))
+    print("R2", r2_score(y_test, y_hat))
+
+    with open(os.path.join(main_folder, model_path), "wb+") as file:
+        pickle.dump(model, file)
+
+    return model
+
+
+def train_model(
+    X_train, X_test, y_train, y_test, columns_target, main_folder, model_path, lr=0.01, seed=SEED, get_history=False
+):
+    # Set all seeds from reproducibility
+    set_all_seeds(seed)
+
+    # Create folder if doesn't exist
+    if not os.path.exists(main_folder):
+        os.makedirs(main_folder)
+
+    ## Setup model for training and training
+    model = setup_model(len(columns_target))
+    opt = tf.keras.optimizers.Adam(learning_rate=lr)  # 0.01 for the hardness
+    print("learning rate", lr)
+    model.compile(optimizer=opt, loss="mean_squared_error")
+
+    validation_split = 0.1
+    history = model.fit(
+        X_train, y_train, batch_size=1, epochs=200, verbose=1, validation_data=(X_test, y_test), shuffle=True
+    )  # 200 epochs initially
+    # raise Exception("Early stopping to test reproducibility")
+    model.save(os.path.join(main_folder, model_path))
+
+    model_core_name = model_path.split(".")[0]
+    with open(os.path.join(main_folder, f"{model_core_name}_fit_history.pickle"), "wb+") as file:
+        pickle.dump(history, file)
+
+    ### Plot loss
+    plt.clf()
+    plt.plot(history.history["loss"])
+    plt.plot(history.history["val_loss"])
+    plt.title("model loss")
+    plt.ylabel("loss")
+    plt.xlabel("epoch")
+    plt.legend(["train", "test"], loc="upper left")
+    fig = plt.gcf()
+    plt.show()
+    fig.savefig(os.path.join(main_folder, "plot_loss_function.png"))
+
+    if get_history:
+        return model, history
+    return model
+
+
+def save_shap_explainer(predict_fn, X_train, X_test, main_folder, explainer_name="explainer"):
+    # Create folder if doesn't exist
+    if not os.path.exists(main_folder):
+        os.makedirs(main_folder)
+
+    ## Get explainer
+    ex = shap.KernelExplainer(predict_fn, X_train[:80])
+    shap_values = ex.shap_values(X_test[-20:])
+    fig, axes = plt.subplots(1, 2, figsize=(5, 5))
+    # need to check that it works in all cases (especially if size the X_test is 1)
+    if len(shap_values) == 1:
+        shap_values = shap_values[0]
+    plt.clf()
+    shap.summary_plot(shap_values, X_test[-20:], show=False)
+    fig = plt.gcf()
+    fig.savefig(os.path.join(main_folder, f"plot_shap_summary_{explainer_name}.png"))
+    plt.show()
+
+    with open(os.path.join(main_folder, f"{explainer_name}.bz2"), "wb") as file:
+        # pickle.dump(ex, file)
+        dill.dump(ex, file)
+
+
+def compute_shap_explainer_no_physics(model_path, X_train, X_test, main_folder, scaler_inputs_path):
+    """
+    Creates and save a shap explainer that do not include physics-informed features
+    To be shared with customers and put into the gradio
+    X_train and X_test must NOT be scaled
+    """
+    scaler_inputs = unpickle_file(scaler_inputs_path)
+    if model_path.split(".")[-1] == "h5":
+        model = tf.keras.models.load_model(model_path)
+    else:
+        model = unpickle_file(model_path)
+
+    model_no_physics = NoPhysicsModels(model, scaler_inputs, add_physics_features)
+
+    save_shap_explainer(model_no_physics.predict, X_train, X_test, main_folder, explainer_name="exp_no_physics")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Process parameters")
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        help="The path to your input data file",
+        default="preprocessed_data.csv",
+        required=False,
+    )
+    parser.add_argument(
+        "--main_folder", type=str, help="Folder to save model files", default="../models/hardness", required=False
+    )
+    parser.add_argument(
+        "--model_path", type=str, help="Path to save model", default="model_hardness.h5", required=False
+    )
+    parser.add_argument("--columns_target", type=str, help="List of target columns", default="H", required=False)
+    parser.add_argument(
+        "--columns_numerical",
+        type=str,
+        help="List of data columns with numeric values",
+        default="%A,%B,%C,%D,%E,%F,%Phase_A,%Phase_B,%Phase_C,%Phase_D,%Phase_E,%Phase_F,%A_Matrice,%B_Matrice,%C_Matrice,%D_Matrice,%E_Matrice,%F_Matrice,H,Temperature_C",
+        required=False,
+    )
+
+    args = parser.parse_args()
+
+    columns_numerical = args.columns_numerical.split(",") if args.columns_numerical else []
+    columns_target = args.columns_target.split(",") if args.columns_target else []
+
+    X_train, X_test, y_train, y_test = prepare_data(args.data_path, columns_numerical, columns_target, args.main_folder)
+
+    model = train_model(X_train, X_test, y_train, y_test, columns_target, args.main_folder, args.model_path)
+
+    save_shap_explainer(model.predict, X_train, X_test, args.main_folder)