dvc_session/src/ensemble.py

import plac
from joblib import load
from sklearn.ensemble import VotingClassifier

from src.utils import (
    dump_yaml,
    evaluate_model,
    print_results,
    read_data,
    read_yaml,
    save_results,
)


@plac.annotations(
    data_path=("Path to source data", "option", "i", str),
    model_path=("Path to save trained Model", "option", "m", str),
    out_path=("Path to save trained Model", "option", "o", str),
)
def main(
    data_path="data/features/",
    model_path="models/",
    out_path="models/ensemble/",
):
    X_train, X_test, y_train, y_test = read_data(data_path)

    name = "Ensemble"
    params = read_yaml("params.yaml", "ensemble")
    cl1 = load(f"{model_path}/logistic/model.gz")
    cl2 = load(f"{model_path}/svc/model.gz")
    cl3 = load(f"{model_path}/r_forrest/model.gz")
    estimators = [("l_regression", cl1), ("l_svc", cl2), ("r_forrest", cl3)]

    model = VotingClassifier(estimators, **params)
    model.fit(X_train, y_train)

    accuracy, c_matrix, fig = evaluate_model(model, X_test, y_test)
    print_results(accuracy, c_matrix, name)

    save_results(out_path, model, fig)
    dump_yaml(
        dict(accuracy=accuracy, confusion_matrix=c_matrix),
        f"{out_path}metrics.yaml",
    )


if __name__ == "__main__":
    plac.call(main)