Vertex AI PipelinesからBigQueryMLで機械学習モデルの訓練を行い予測を実行する

PROJECT_ID = "自分のGoogleCloudプロジェクトID"
BQ_DATASET = "kubeflow_dataset"
RAW_DATA_TABLE = "penguins"
PREPROCESSED_TABLE = "penguins_processed"
CLASSIFICATION_MODEL_NAME = "penguins_model"
PREDICT_RESULT_TABLE = "penguins_predict_result"


def get_project():
    """ プロジェクトIDを返す
    """
    return PROJECT_ID


def get_predict_job_config():
    """ 推定結果を格納するテーブルの情報をジョブ設定として返す
    """
    job_config = {
        "destinationTable": {
            "projectId": PROJECT_ID,
            "datasetId": BQ_DATASET,
            "tableId": PREDICT_RESULT_TABLE,
        }
    }
    return job_config


def generate_create_bq_preprocess_query():
    """ 前処理のためのSQLを返す
    """
    create_bq_preprocess_query = f"""
        -- create the splitted table
        CREATE OR REPLACE TABLE `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}`
        AS (
            WITH data_with_split_flag AS
            (
                SELECT
                    species,
                    island,
                    culmen_length_mm,
                    culmen_depth_mm,
                    flipper_length_mm,
                    body_mass_g,
                    sex,
                    RAND() AS split_flag
                FROM `{PROJECT_ID}.{BQ_DATASET}.{RAW_DATA_TABLE}`
                WHERE
                    body_mass_g IS NOT NULL
            )
            SELECT
                species,
                island,
                culmen_length_mm,
                culmen_depth_mm,
                flipper_length_mm,
                body_mass_g,
                sex,
                CASE WHEN split_flag < 0.8 THEN 'TRAIN' ELSE 'PREDICT' END AS split
            FROM data_with_split_flag
          )
        """
    return create_bq_preprocess_query


def generate_create_bq_model_query():
    """ モデル作成を実行するSQLを返す
    """
    create_bq_model_query = f"""
        CREATE OR REPLACE MODEL `{PROJECT_ID}.{BQ_DATASET}.{CLASSIFICATION_MODEL_NAME}`
          OPTIONS (
              model_type='linear_reg',
              input_label_cols=['body_mass_g']) AS
          SELECT
              species,
              island,
              culmen_length_mm,
              culmen_depth_mm,
              flipper_length_mm,
              body_mass_g,
              sex
          FROM
             `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}`
          WHERE split = 'TRAIN';
        """
    return create_bq_model_query


def generate_create_bq_prediction_query():
    """ モデルで推定する対象のデータを表すSQLを返す
    """
    create_bq_prediction_query = f"""
            SELECT
                species,
                island,
                culmen_length_mm,
                culmen_depth_mm,
                flipper_length_mm,
                sex
            FROM `{PROJECT_ID}.{BQ_DATASET}.{PREPROCESSED_TABLE}` 
            WHERE split = 'PREDICT' 
        """
    return create_bq_prediction_query