stackexchange-tutorial/tutorial/train_model.py

import pandas as pd
from . import shared


def fit_model(params: dict):
    tfidf = shared.load(shared.train_tfidf)
    y = shared.load_labels(shared.train_data)

    from sklearn.dummy import DummyClassifier as Classifier
    clf = Classifier(strategy='stratified', **params)
    clf.fit(tfidf, y)

    return y, clf.predict(tfidf), clf


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


def main(params: dict):
    y, preds, clf = fit_model(params)
    metrics = eval_on_train_data(y, preds)
    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 y, preds, 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)