Cookiecutter-MLOps/src/models/train_model.py

import sys
import os

import pandas as pd
from scipy.sparse import load_npz
import pickle
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score
from dagshub import DAGsHubLogger
from tqdm import tqdm


def train(processed_data_path, model_folder_path):
    X_train = load_npz(processed_data_path + "X_train.npz")
    y_train = pd.read_csv(processed_data_path + "y_train.csv")

    logger = DAGsHubLogger(
        metrics_path="reports/training_metrics.csv",
        hparams_path="src/models/training_params.yml",
    )

    val_error_rate = []
    neighbors_range = range(1, 500, 5)

    for i in tqdm(neighbors_range):
        knn = KNeighborsClassifier(n_neighbors=i)
        val_error = (
            1 - cross_val_score(knn, X_train, y_train.values.ravel(), cv=2).mean()
        )
        logger.log_metrics({"val_error": val_error}, step_num=i)
        val_error_rate.append(val_error)

    best_k = neighbors_range[val_error_rate.index(min(val_error_rate))]
    logger.log_hyperparams(best_k=best_k)

    knn = KNeighborsClassifier(n_neighbors=best_k)
    knn.fit(X_train, y_train.values.ravel())

    # source, destination
    pickle.dump(knn, open(model_folder_path + "model.pkl", "wb"))

    logger.save()
    logger.close()


if __name__ == "__main__":
    if not (1 <= len(sys.argv) <= 3):
        print(
            "usage: %s <processed_data_folder (optional)> <out_folder (optional)>"
            % sys.argv[0],
            file=sys.stderr,
        )
        sys.exit(0)

    in_folder = sys.argv[1] if len(sys.argv) >= 2 else "data/processed/"

    out_folder = sys.argv[2] if len(sys.argv) == 3 else "models/"

    if not os.path.exists(out_folder):
        os.makedirs(out_folder)

    train(in_folder, out_folder)