dvc_session/src/random_forrest.py

import plac
from sklearn.ensemble import RandomForestClassifier

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),
    n_estimators=("Path to save trained Model", "option", "e", str),
    max_samples=("Path to save trained Model", "option", "s", str),
    out_path=("Path to save trained Model", "option", "o", str),
)
def main(
    data_path="data/features/",
    out_path="models/r_forrest/",
    n_estimators=10,
    max_samples=30,
):
    X_train, X_test, y_train, y_test = read_data(data_path)

    name = "RandomForrest"
    params = read_yaml("params.yaml", "forrest")
    model = RandomForestClassifier(**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)