""" Pre Release Forecasting Pipleines """ from forecasting_toolkit.models.model_factory.baseline_models import ( linear_regressor_pipeline_factory, gbt_regressor_pipeline_factory, histgbm_regressor_pipeline_factory ) from forecasting_toolkit.feature_preprocessing.timeseries.pipelines import ( feature_preprocessing_pipeline_factory_v2 ) from forecasting_toolkit.model_evaluation.timeseries.pipelines import ( evaluate_generalization_performance_sklearn ) from forecasting_toolkit.datastore.adapters.helpers import ( fetch_dataset ) from forecasting_toolkit.utils import ( has_columns, save_sklearn_model ) # std libs import os import argparse # misc import numpy as np from absl import logging import mlflow # sklearn imports from sklearn.model_selection import TimeSeriesSplit from sklearn.pipeline import Pipeline # forecasting toolkit imports from forecasting_toolkit.utils import ( check_mlflow_connection ) # set tracking uri TRACKING_SERVER = os.environ.get("TRACK_URI", "https://dev-orch-mlflow-service.dev.theorchard.io/") mlflow.set_tracking_uri(TRACKING_SERVER) """ Modelling Pipelines """ def linear_regressor_baseline_v2(dataset_table="DATASET_STREAMS_DAILY_2022", store_id=286, target_col="STREAMS", model_params={}): """ Linear Regressor (v2) Trains a baseline linear regressor pipeline params: dataset_table (str) - Snowflake table that has the dataset store_id (int) - Store id model_params (dict) - model dictionary returns: trained model (sklearn.Pipeline) - sklearn model trained model """ # fetch dataset dataset_df = fetch_dataset(snowflake_table=dataset_table, filters={"STORE_ID": store_id}) # define feature vector params numerical_cols = [ 'TOTAL_WEEKLY_STREAMS_ROLLUP', 'AVG_WEEKLY_STREAMS_ROLLUP', 'MIN_WEEKLY_STREAMS_ROLLUP', 'MAX_WEEKLY_STREAMS_ROLLUP', 'TOTAL_WEEKLY_SKIPS_ROLLUP', 'AVG_WEEKLY_SKIPS_ROLLUP', 'MIN_WEEKLY_SKIPS_ROLLUP', 'MAX_WEEKLY_SKIPS_ROLLUP', 'TOTAL_WEEKLY_SAVES_ROLLUP', 'AVG_WEEKLY_SAVES_ROLLUP', 'MIN_WEEKLY_SAVES_ROLLUP', 'MAX_WEEKLY_SAVES_ROLLUP', 'TOTAL_WEEKLY_STREAMS_PASSIVE_ROLLUP','AVG_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'AVG_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'AVG_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MIN_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MAX_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ROLLUP','AVG_MONTHLY_STREAMS_ROLLUP', 'MIN_MONTHLY_STREAMS_ROLLUP', 'MAX_MONTHLY_STREAMS_ROLLUP', 'TOTAL_MONTHLY_SKIPS_ROLLUP', 'AVG_MONTHLY_SKIPS_ROLLUP', 'MIN_MONTHLY_SKIPS_ROLLUP', 'MAX_MONTHLY_SKIPS_ROLLUP', 'TOTAL_MONTHLY_SAVES_ROLLUP', 'AVG_MONTHLY_SAVES_ROLLUP', 'MIN_MONTHLY_SAVES_ROLLUP', 'MAX_MONTHLY_SAVES_ROLLUP', 'TOTAL_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'AVG_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MIN_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MAX_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'DURATION'] time_cols = [ 'SNAPSHOT_ISOWEEK', 'SNAPSHOT_YEAR', 'SNAPSHOT_MONTH', 'SNAPSHOT_DAY_OF_YEAR', 'SNAPSHOT_DAY_OF_WEEK', 'RELEASE_YEAR', 'RELEASE_MONTH', 'RELEASE_DOW', 'RELEASE_WEEKISO', 'RELEASE_AGE_WEEKS', 'RELEASE_AGE_DAYS'] # ensure we have the columns we need try: logging.debug('checing columns') assert has_columns(columns=numerical_cols, dataframe=dataset_df) assert has_columns(columns=time_cols, dataframe=dataset_df) except AssertionError: logging.erorr("Missing Columns in dataset_df") # feature dictionary feature_dict = { "numerical_float_cols": numerical_cols, "time_cols": time_cols, "bool_cols": ['THIRD_PARTY_PUBLISHER'], "snapshot_year": 'SNAPSHOT_YEAR', "snapshot_day_of_week_col": 'SNAPSHOT_DAY_OF_WEEK', "snapshot_day_of_year_col": 'SNAPSHOT_DAY_OF_YEAR', "snapshot_month_col": 'SNAPSHOT_MONTH', "release_year": 'RELEASE_YEAR', "release_month_col": 'RELEASE_MONTH', "release_week_col": 'RELEASE_WEEKISO', "release_day_of_week_col": 'RELEASE_DOW', "categorical_cols": ['ARTIST_ID', 'LABEL_ID', 'RELEASE_GENREID', 'TRACK_TYPE'] } # additional mlflow experiment tags mlflow_tags = { "pipeline": __name__, "team": os.environ.get("TEAM", "data-platform"), "feature_processing_pipeline": "feature_preprocessing_pipeline_factory_v2", "feature_dict": feature_dict } logging.debug("mlflow tags") logging.debug(mlflow_tags) # feature engineering pipeline logging.debug("generating feature preprocessing pipeline") feat_pipeline_v2 = feature_preprocessing_pipeline_factory_v2(**feature_dict) # transform dataset logging.debug("tranforming dataset") X_feat = feat_pipeline_v2.fit_transform(dataset_df) y = np.array(dataset_df[target_col].values) # evaluate model performance logging.debug("evaluating model on dataset") timseries_cv = TimeSeriesSplit(n_splits=5, test_size=7, gap=0) evaluate_generalization_performance_sklearn(model_factory=linear_regressor_pipeline_factory, X=X_feat, y=y, cv=timseries_cv, model_params=model_params) # create linear regressor pipleine logging.debug("creating linear regressor ") linear_regressor = linear_regressor_pipeline_factory(model_params=model_params) # fit on dataset and return model logging.debug("fitting model") model = linear_regressor.fit(X_feat, y) return model def gbt_regressor_baseline_v2(dataset_table:str="DATASET_STREAMS_DAILY_2022", store_id:int=286, country_code:str="US", fit_only:bool=False, target_col:str= "STREAMS", model_params:dict={'n_estimators': 100, 'learning_rate': 0.1, 'subsample': 1.0, 'min_samples_split':3, 'min_samples_leaf': 10, 'max_depth': 3, 'validation_fraction': 0.3, 'n_iter_no_change': 5}) -> Pipeline: """ Gradient Boosted Tree Regressor (v2) Using Feature engineering pipeline v2 """ # fetch dataset dataset_df = fetch_dataset(snowflake_table=dataset_table, filters={"STORE_ID": store_id, "COUNTRY_CODE": country_code }) # define feature vector params numerical_cols = [ 'TOTAL_WEEKLY_STREAMS_ROLLUP', 'AVG_WEEKLY_STREAMS_ROLLUP', 'MIN_WEEKLY_STREAMS_ROLLUP', 'MAX_WEEKLY_STREAMS_ROLLUP', 'TOTAL_WEEKLY_SKIPS_ROLLUP', 'AVG_WEEKLY_SKIPS_ROLLUP', 'MIN_WEEKLY_SKIPS_ROLLUP', 'MAX_WEEKLY_SKIPS_ROLLUP', 'TOTAL_WEEKLY_SAVES_ROLLUP', 'AVG_WEEKLY_SAVES_ROLLUP', 'MIN_WEEKLY_SAVES_ROLLUP', 'MAX_WEEKLY_SAVES_ROLLUP', 'TOTAL_WEEKLY_STREAMS_PASSIVE_ROLLUP','AVG_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'AVG_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'AVG_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MIN_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MAX_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ROLLUP','AVG_MONTHLY_STREAMS_ROLLUP', 'MIN_MONTHLY_STREAMS_ROLLUP', 'MAX_MONTHLY_STREAMS_ROLLUP', 'TOTAL_MONTHLY_SKIPS_ROLLUP', 'AVG_MONTHLY_SKIPS_ROLLUP', 'MIN_MONTHLY_SKIPS_ROLLUP', 'MAX_MONTHLY_SKIPS_ROLLUP', 'TOTAL_MONTHLY_SAVES_ROLLUP', 'AVG_MONTHLY_SAVES_ROLLUP', 'MIN_MONTHLY_SAVES_ROLLUP', 'MAX_MONTHLY_SAVES_ROLLUP', 'TOTAL_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'AVG_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MIN_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MAX_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'DURATION'] time_cols = [ 'SNAPSHOT_ISOWEEK', 'SNAPSHOT_YEAR', 'SNAPSHOT_MONTH', 'SNAPSHOT_DAY_OF_YEAR', 'SNAPSHOT_DAY_OF_WEEK', 'RELEASE_YEAR', 'RELEASE_MONTH', 'RELEASE_DOW', 'RELEASE_WEEKISO', 'RELEASE_AGE_WEEKS', 'RELEASE_AGE_DAYS'] # ensure we have the columns we need try: logging.debug('checing columns') assert has_columns(columns=numerical_cols, dataframe=dataset_df) == True assert has_columns(columns=time_cols, dataframe=dataset_df) == True except AssertionError: logging.erorr("Missing Columns in dataset_df") # feature dictionary feature_dict = { "numerical_float_cols":numerical_cols, "time_cols": time_cols, "bool_cols":['THIRD_PARTY_PUBLISHER'], "snapshot_year": 'SNAPSHOT_YEAR', "snapshot_day_of_week_col": 'SNAPSHOT_DAY_OF_WEEK', "snapshot_day_of_year_col": 'SNAPSHOT_DAY_OF_YEAR', "snapshot_month_col": 'SNAPSHOT_MONTH', "release_year": 'RELEASE_YEAR', "release_month_col": 'RELEASE_MONTH', "release_week_col": 'RELEASE_WEEKISO', "release_day_of_week_col": 'RELEASE_DOW', "categorical_cols":['ARTIST_ID', 'LABEL_ID', 'RELEASE_GENREID', 'TRACK_TYPE'] } # additional mlflow experiment tags mlflow_tags = { "pipeline": __name__, "team": os.environ.get("TEAM", "data-platform"), "feature_processing_pipeline": "feature_preprocessing_pipeline_factory_v2", "feature_dict": feature_dict, "store_id": store_id, "country_code": country_code } logging.debug("mlflow tags") logging.debug(mlflow_tags) # feature processing pipeline logging.debug("creating feature preprocessing pipeline") feat_pipeline_v2 = feature_preprocessing_pipeline_factory_v2(**feature_dict) # transform features logging.debug("transforming dataset with feature preprocessing pipeline") y = np.array(dataset_df[target_col].values) if not fit_only: # fit transform X_feat = feat_pipeline_v2.fit_transform(dataset_df) # evaluate model performance logging.debug("evaluating model performance") evaluate_generalization_performance_sklearn(model_factory=gbt_regressor_pipeline_factory, X=X_feat, y=y, model_params=model_params, mlflow_experiment=f"data-platform/debut-forecasting/{country_code}/{store_id}") # create fresh gbt regressor pipleine logging.debug('initialising gbt regressor') gbt_regressor = gbt_regressor_pipeline_factory(model_params=model_params) gbt_pipeline = Pipeline(steps=[("feat_pipeline_v2", feat_pipeline_v2), ("gbt_regressor", gbt_regressor)]) # fit on dataset and return model logging.debug('fitting gbt regressor') fitted_model = gbt_pipeline.fit(dataset_df, y) return fitted_model def histgbm_regressor_baseline_v1(dataset_table:str="DATASET_STREAMS_DAILY_2022", store_id:int=286, country_code:str="US", fit_only:bool=False, target_col:str= "STREAMS", model_params:dict={}) -> Pipeline: """ LightGBM Using Feature engineering pipeline v2 params: dataset_table (str) -snowflake table containing the dataset store_id (int) - store id country_code (str) - country code to fit the dataset on fit_only (bool) - only fits the model on the dataset without evaluating expected generalization performance of the model target_col (str) - The column we want our model to learn to predict model_params (dict) - Dictionary containing customised model parameters """ # fetch dataset dataset_df = fetch_dataset(snowflake_table=dataset_table, filters={"STORE_ID": store_id, "COUNTRY_CODE": country_code }) # define feature vector params numerical_cols = [ 'TOTAL_WEEKLY_STREAMS_ROLLUP', 'AVG_WEEKLY_STREAMS_ROLLUP', 'MIN_WEEKLY_STREAMS_ROLLUP', 'MAX_WEEKLY_STREAMS_ROLLUP', 'TOTAL_WEEKLY_SKIPS_ROLLUP', 'AVG_WEEKLY_SKIPS_ROLLUP', 'MIN_WEEKLY_SKIPS_ROLLUP', 'MAX_WEEKLY_SKIPS_ROLLUP', 'TOTAL_WEEKLY_SAVES_ROLLUP', 'AVG_WEEKLY_SAVES_ROLLUP', 'MIN_WEEKLY_SAVES_ROLLUP', 'MAX_WEEKLY_SAVES_ROLLUP', 'TOTAL_WEEKLY_STREAMS_PASSIVE_ROLLUP','AVG_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'AVG_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MIN_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'MAX_WEEKLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'AVG_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MIN_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'MAX_WEEKLY_STREAMS_COLLECTION_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ROLLUP','AVG_MONTHLY_STREAMS_ROLLUP', 'MIN_MONTHLY_STREAMS_ROLLUP', 'MAX_MONTHLY_STREAMS_ROLLUP', 'TOTAL_MONTHLY_SKIPS_ROLLUP', 'AVG_MONTHLY_SKIPS_ROLLUP', 'MIN_MONTHLY_SKIPS_ROLLUP', 'MAX_MONTHLY_SKIPS_ROLLUP', 'TOTAL_MONTHLY_SAVES_ROLLUP', 'AVG_MONTHLY_SAVES_ROLLUP', 'MIN_MONTHLY_SAVES_ROLLUP', 'MAX_MONTHLY_SAVES_ROLLUP', 'TOTAL_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_PASSIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'AVG_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MIN_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'MAX_MONTHLY_STREAMS_ACTIVE_ROLLUP', 'TOTAL_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'AVG_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MIN_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'MAX_MONTHLY_STREAMS_COLLECTION_ROLLUP', 'DURATION'] time_cols = [ 'SNAPSHOT_ISOWEEK', 'SNAPSHOT_YEAR', 'SNAPSHOT_MONTH', 'SNAPSHOT_DAY_OF_YEAR', 'SNAPSHOT_DAY_OF_WEEK', 'RELEASE_YEAR', 'RELEASE_MONTH', 'RELEASE_DOW', 'RELEASE_WEEKISO', 'RELEASE_AGE_WEEKS', 'RELEASE_AGE_DAYS'] # ensure we have the columns we need try: logging.debug('checing columns') assert has_columns(columns=numerical_cols, dataframe=dataset_df) == True assert has_columns(columns=time_cols, dataframe=dataset_df) == True except AssertionError: logging.erorr("Missing Columns in dataset_df") # feature dictionary feature_dict = { "numerical_float_cols":numerical_cols, "time_cols": time_cols, "bool_cols":['THIRD_PARTY_PUBLISHER'], "snapshot_year": 'SNAPSHOT_YEAR', "snapshot_day_of_week_col": 'SNAPSHOT_DAY_OF_WEEK', "snapshot_day_of_year_col": 'SNAPSHOT_DAY_OF_YEAR', "snapshot_month_col": 'SNAPSHOT_MONTH', "release_year": 'RELEASE_YEAR', "release_month_col": 'RELEASE_MONTH', "release_week_col": 'RELEASE_WEEKISO', "release_day_of_week_col": 'RELEASE_DOW', "categorical_cols":['ARTIST_ID', 'LABEL_ID', 'RELEASE_GENREID', 'TRACK_TYPE'] } # additional mlflow experiment tags mlflow_tags = { "pipeline": __name__, "team": os.environ.get("TEAM", "data-platform"), "feature_processing_pipeline": "feature_preprocessing_pipeline_factory_v2", "feature_dict": feature_dict, "store_id": store_id, "country_code": country_code } logging.debug("mlflow tags") logging.debug(mlflow_tags) # feature processing pipeline logging.debug("creating feature preprocessing pipeline") feat_pipeline_v2 = feature_preprocessing_pipeline_factory_v2(**feature_dict) # transform features logging.debug("transforming dataset with feature preprocessing pipeline") y = np.array(dataset_df[target_col].values) if not fit_only: # fit transform X_feat = feat_pipeline_v2.fit_transform(dataset_df) # convert to dense dataset X_feat = X_feat.toarray() # evaluate model performance logging.debug("evaluating model performance") evaluate_generalization_performance_sklearn(model_factory=histgbm_regressor_pipeline_factory, X=X_feat, y=y, model_params=model_params, mlflow_experiment=f"data-platform/debut-forecasting/{country_code}/{store_id}") # create fresh gbt regressor pipleine logging.debug('initialising gbt regressor') hist_gbm_regressor = histgbm_regressor_pipeline_factory(model_params=model_params) gbt_pipeline = Pipeline(steps=[("feat_pipeline_v2", feat_pipeline_v2), ("gbt_regressor", hist_gbm_regressor) ]) # fit on dataset and return model logging.debug('fitting gbt regressor') fitted_model = gbt_pipeline.fit(dataset_df, y) return fitted_model if __name__ == '__main__': parser = argparse.ArgumentParser(description='Fits a model to the given dataset and saves a model') parser.add_argument('--dataset', type=str, help='Store ID', dest="snowflake_dataset_table", default="DATASET_STREAMS_DAILY_2022") parser.add_argument('--store_id', type=int, help='Store ID', dest="store_id") parser.add_argument('--country_code', type=str, help='Country Code', default="US", dest="country_code") parser.add_argument('--model', type=str, choices=["gbt_v2", "histgbm_v1"], default="gbt_v2", help='Model used to fit the dataset', dest="selected_model", required=False) parser.add_argument('--output', type=str, default="./model.pkl", help='Model Output file - Saves model locally', dest="model_output_path", required=False) parser.add_argument('--fit_only', type=bool, default=False, help='Fits the model only without evaluation', dest="fit_only", required=False) parser.add_argument('--initial_model', help='Warm start from an initial model - mlflow model ID', dest="selected_model", required=False) parser.add_argument('--model_name', help='Name to register the model with mlflow service', dest="model_name", default="gbt_model", required=False) parser.add_argument('--debug', default=True, help='Sets logging to debug mode', dest="debug", required=False) # parse arguments args = parser.parse_args() if args.debug: logging.set_verbosity(logging.DEBUG) # check tracking server connection check_mlflow_connection(TRACKING_SERVER) # Save model model_output_path = args.model_output_path if args.selected_model in ["gbt_v2"]: # evaluate and fit gbt regressor logging.debug("Fitting GBT Regressor Baseline - V2") model = gbt_regressor_baseline_v2(store_id=args.store_id, fit_only=args.fit_only, country_code=args.country_code) logging.debug('Saving GBT Regressor Basline -v2') save_sklearn_model(model=model, model_path=model_output_path) elif args.selected_model in ["histgbm_v1"]: # evaluate and fit gbt regressor (very similar to lightgbm implementation wise) logging.debug("Fitting GBT Regressor Baseline - V2") model = histgbm_regressor_baseline_v1(store_id=args.store_id, fit_only=args.fit_only, country_code=args.country_code) logging.debug('Saving Histogram GBT Regressor Basline -v2') save_sklearn_model(model=model, model_path=model_output_path)