autogen/test/pipeline_tuning_example/train/train.py

import argparse
import lightgbm as lgb
import os
import pandas as pd
from azureml.core import Run


class LightGBMCallbackHandler():
    def __init__(self):
        pass

    def callback(self, env: lgb.callback.CallbackEnv) -> None:
        """Callback method to collect metrics produced by LightGBM.

        See https://lightgbm.readthedocs.io/en/latest/_modules/lightgbm/callback.html
        """
        # loop on all the evaluation results tuples
        print("env.evaluation_result_list:", env.evaluation_result_list)
        for data_name, eval_name, result, _ in env.evaluation_result_list:
            run = Run.get_context()
            run.log(f"{data_name}_{eval_name}", result)


def main(args):
    """Main function of the script."""

    train_path = os.path.join(args.train_data, 'data.csv')
    print("traning_path:", train_path)

    test_path = os.path.join(args.test_data, 'data.csv')

    train_set = lgb.Dataset(train_path)
    test_set = lgb.Dataset(test_path)
    callbacks_handler = LightGBMCallbackHandler()
    config = {"header": True, "objective": "binary", "label_column": 30, "metric": "binary_error",
              "n_estimators": args.n_estimators, "learning_rate": args.learning_rate}
    gbm = lgb.train(
        config,
        train_set,
        valid_sets=[test_set],
        valid_names=["eval"],
        callbacks=[
            callbacks_handler.callback,
        ],
    )

    print('Saving model...')
    # save model to file
    gbm.save_model(os.path.join(args.model, 'model.txt'))


if __name__ == "__main__":
    # input and output arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_data", type=str, help="path to train data")
    parser.add_argument("--test_data", type=str, help="path to test data")
    parser.add_argument("--n_estimators", required=False, default=100, type=int)
    parser.add_argument("--learning_rate", required=False, default=0.1, type=float)
    parser.add_argument("--model", type=str, help="path to output directory")
    args = parser.parse_args()
    main(args)