stackexchange-tutorial/tutorial/train_model.py

import pandas as pd
from . import shared


def fit_model(params: dict):
    print("Loading training data")
    X, y = shared.load_data_and_labels(shared.train_processed)
    print("Done")

    from sklearn.linear_model import LogisticRegression as Classifier
    clf = Classifier(**params)
    print("Training model ", clf)
    # Required for efficient training, so that sklearn doesn't inflate the pandas sparse DF to a dense matrix.
    # sklearn only supports scipy sparse matrices.
    X_sparse = X.sparse.to_coo() 
    clf.fit(X_sparse, y)
    print("Done")

    return X, y, clf


def eval_on_train_data(clf, X, y):
    return shared.compute_metrics(clf, X, y, "train")


def main(params: dict):
    X, y, clf = fit_model(params)
    metrics = eval_on_train_data(clf, X, y)
    from dagshub import dagshub_logger
    with dagshub_logger() as logger:
        logger.log_hyperparams(clf.get_params(), classifier_type=type(clf).__name__)
        logger.log_metrics(metrics)
    shared.save(clf, shared.classifier_pkl)
    
    # For possible interactive use
    return X, y, clf, metrics


if __name__ == "__main__":
    def param_key(arg):
        prefix = '--param-'
        if arg.startswith(prefix):
            arg = arg[len(prefix):]
        arg = arg.replace('-','_')
        return arg

    import sys
    params = {param_key(k):v for k,v in zip(sys.argv[1::2],sys.argv[2::2])}
    print('Running training with params: ', params)
    main(params)