# Copyright 2021 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

name: tabular_stats_and_example_gen
description: |
  Statistics and example gen for tabular data

    Args:
        project (str):
            Required. Project to run dataset statistics and example generation.
        location (str):
            Location for running dataset statistics and example generation.
        root_dir (str):
            The Cloud Storage location to store the output.
        target_column_name (str):
            The target column name.
        weight_column_name (str):
            The weight column name.
        prediction_type (str):
            The prediction type. Supported values: "classification", "regression".
        optimization_objective (str):
            Objective function the model is optimizing towards. The training process creates a model
            that maximizes/minimizes the value of the objective function over the validation set.
            The supported optimization objectives depend on the prediction type. If the field is not
            set, a default objective function is used.
            classification (binary):
                "maximize-au-roc" (default) - Maximize the area under the receiver
                    operating characteristic (ROC) curve.
                "minimize-log-loss" - Minimize log loss.
                "maximize-au-prc" - Maximize the area under the precision-recall curve.
                "maximize-precision-at-recall" - Maximize precision for a specified recall value.
                "maximize-recall-at-precision" - Maximize recall for a specified precision value.
            classification (multi-class):
                "minimize-log-loss" (default) - Minimize log loss.
            regression:
                "minimize-rmse" (default) - Minimize root-mean-squared error (RMSE).
                "minimize-mae" - Minimize mean-absolute error (MAE).
                "minimize-rmsle" - Minimize root-mean-squared log error (RMSLE).
        optimization_objective_recall_value (str):
            Required when optimization_objective is "maximize-precision-at-recall". Must be between
            0 and 1, inclusive.
        optimization_objective_precision_value (str):
            Required when optimization_objective is "maximize-recall-at-precision". Must be between
            0 and 1, inclusive.
        transformations (str):
            Quote escaped JSON string for transformations. Each transformation will apply transform
            function to given input column. And the result will be used for training. When creating
            transformation for BigQuery Struct column, the column should be flattened using "." as
            the delimiter.
        transformations_path (Optional[str]):
            Path to a GCS file containing JSON string for transformations.
        dataflow_machine_type (Optional[str]):
            The machine type used for dataflow jobs. If not set, default to n1-standard-16.
        dataflow_max_num_workers (Optional[int]):
            The number of workers to run the dataflow job. If not set, default to 25.
        dataflow_disk_size_gb (Optional[int]):
            The disk size, in gigabytes, to use on each Dataflow worker instance. If not set,
            default to 40.
        dataflow_subnetwork (Optional[str]):
            Dataflow's fully qualified subnetwork name, when empty the default subnetwork will be
            used. More details:
            https://cloud.google.com/dataflow/docs/guides/specifying-networks#example_network_and_subnetwork_specifications
        dataflow_use_public_ips (Optional[bool]):
            Specifies whether Dataflow workers use public IP addresses.
        dataflow_service_account (Optional[str]):
            Custom service account to run dataflow jobs.
        encryption_spec_key_name (Optional[str]):
            Customer-managed encryption key.
        run_distillation (bool): True if in distillation mode. The default value is false.

    Returns:
        dataset_schema (DatasetSchema):
            The schema of the dataset.
        dataset_stats (AutoMLTabularDatasetStats):
            The stats of the dataset.
        train_split (Dataset):
            The train split.
        eval_split (Dataset):
            The eval split.
        test_split (Dataset):
            The test split.
        test_split_json (JsonObject):
            The test split JSON object.
        downsampled_test_split_json (JsonObject):
            The downsampled test split JSON object.
        instance_baseline (AutoMLTabularInstanceBaseline):
            The instance baseline used to calculate explanations.
        metadata (TabularExampleGenMetadata):
            The tabular example gen metadata.
        gcp_resources (str):
            GCP resources created by this component.
            For more details, see https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/proto/README.md.
inputs:
- {name: project, type: String}
- {name: location, type: String}
- {name: root_dir, type: String}
- {name: target_column_name, type: String}
- {name: weight_column_name, type: String, default: ""}
- {name: prediction_type, type: String}
- {name: optimization_objective, type: String, default: ""}
- {name: optimization_objective_recall_value, type: Float, default: "-1"}
- {name: optimization_objective_precision_value, type: Float, default: "-1"}
- {name: transformations, type: String}
- {name: transformations_path, type: String, default: ""}
- {name: request_type, type: String, default: "COLUMN_STATS_ONLY"}
- {name: dataflow_machine_type, type: String, default: "n1-standard-16"}
- {name: dataflow_max_num_workers, type: Integer, default: "25"}
- {name: dataflow_disk_size_gb, type: Integer, default: "40"}
- {name: dataflow_subnetwork, type: String, default: ""}
- {name: dataflow_use_public_ips, type: Boolean, default: "true"}
- {name: dataflow_service_account, type: String, default: ""}
- {name: encryption_spec_key_name, type: String, default: ""}
- {name: run_distillation, type: Boolean, default: "false"}
- {name: additional_experiments, type: String, default: ""}
- {name: additional_experiments_json, type: JsonObject, default: "{}"}
- {name: data_source_csv_filenames, type: String, default: ""}
- {name: data_source_bigquery_table_path, type: String, default: ""}
- {name: predefined_split_key, type: String, default: ""}
- {name: timestamp_split_key, type: String, default: ""}
- {name: stratified_split_key, type: String, default: ""}
- {name: training_fraction, type: Float, default: "-1"}
- {name: validation_fraction, type: Float, default: "-1"}
- {name: test_fraction, type: Float, default: "-1"}
- {name: quantiles, type: JsonArray, default: "[]"}
- {name: enable_probabilistic_inference, type: Boolean, default: "false"}

outputs:
- {name: dataset_schema, type: DatasetSchema}
- {name: dataset_stats, type: AutoMLTabularDatasetStats}
- {name: train_split, type: Dataset}
- {name: eval_split, type: Dataset}
- {name: test_split, type: Dataset}
- {name: test_split_json, type: JsonArray}
- {name: downsampled_test_split_json, type: JsonArray}
- {name: instance_baseline, type: AutoMLTabularInstanceBaseline}
- {name: metadata, type: TabularExampleGenMetadata}
- {name: gcp_resources, type: String}

implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:1.0.32
    command: [python3, -u, -m, google_cloud_pipeline_components.container.v1.custom_job.launcher]
    args: [
      --type, CustomJob,
      --project, {inputValue: project},
      --location, {inputValue: location},
      --gcp_resources, {outputPath: gcp_resources},
      --payload,
      concat: [
        '{"display_name": "tabular-stats-and-example-gen-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}", "encryption_spec": {"kms_key_name":"',
        {inputValue: encryption_spec_key_name},
        '"}, "job_spec": {"worker_pool_specs": [{"replica_count": 1, "machine_spec": {"machine_type": "n1-standard-8"}, "container_spec": {"image_uri":"',
        'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/training:20230322_0325',
        '", "args": ["stats_generator",',
        '"--train_spec={\"prediction_type\": \"',
        {inputValue: prediction_type},
        '\", \"target_column\": \"',
        {inputValue: target_column_name},
        '\", \"optimization_objective\": \"',
        {inputValue: optimization_objective},
        '\", \"weight_column_name\": \"',
        {inputValue: weight_column_name},
        '\", \"transformations\": ',
        {inputValue: transformations},
        ', \"quantiles\": ',
        {inputValue: quantiles},
        ', \"enable_probabilistic_inference\": ',
        {inputValue: enable_probabilistic_inference},
        '}", "--transformations_override_path=',
        {inputValue: transformations_path},
        '", "--data_source_csv_filenames=',
        {inputValue: data_source_csv_filenames},
        '", "--data_source_bigquery_table_path=',
        {inputValue: data_source_bigquery_table_path},
        '", "--predefined_split_key=',
        {inputValue: predefined_split_key},
        '", "--timestamp_split_key=',
        {inputValue: timestamp_split_key},
        '", "--stratified_split_key=',
        {inputValue: stratified_split_key},
        '", "--training_fraction=',
        {inputValue: training_fraction},
        '", "--validation_fraction=',
        {inputValue: validation_fraction},
        '", "--test_fraction=',
        {inputValue: test_fraction},
        '", "--target_column=',
        {inputValue: target_column_name},
        '", "--request_type=',
        {inputValue: request_type},
        '", "--optimization_objective_recall_value=',
        {inputValue: optimization_objective_recall_value},
        '", "--optimization_objective_precision_value=',
        {inputValue: optimization_objective_precision_value},
        '", "--example_gen_gcs_output_prefix=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/example_gen_output", "--dataset_stats_dir=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/stats/", "--stats_result_path=',
        {outputUri: dataset_stats},
        '", "--dataset_schema_path=',
        {outputUri: dataset_schema},
        '", "--job_name=tabular-stats-and-example-gen-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}',
        '", "--dataflow_project=',
        {inputValue: project},
        '", "--error_file_path=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/error.pb", "--dataflow_staging_dir=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/dataflow_staging", "--dataflow_tmp_dir=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/dataflow_tmp", "--dataflow_max_num_workers=',
        {inputValue: dataflow_max_num_workers},
        '", "--dataflow_worker_container_image=',
        'us-docker.pkg.dev/vertex-ai/automl-tabular/dataflow-worker:20230322_0325',
        '", "--dataflow_machine_type=',
        {inputValue: dataflow_machine_type},
        '", "--dataflow_disk_size_gb=',
        {inputValue: dataflow_disk_size_gb},
        '", "--dataflow_kms_key=',
        {inputValue: encryption_spec_key_name},
        '", "--dataflow_subnetwork_fully_qualified=',
        {inputValue: dataflow_subnetwork},
        '", "--dataflow_use_public_ips=',
        {inputValue: dataflow_use_public_ips},
        '", "--dataflow_service_account=',
        {inputValue: dataflow_service_account},
        '", "--is_distill=',
        {inputValue: run_distillation},
         '", "--additional_experiments=',
        {inputValue: additional_experiments},
        '", "--metadata_path=',
        {outputUri: metadata},
        '", "--train_split=',
        {outputUri: train_split},
        '", "--eval_split=',
        {outputUri: eval_split},
        '", "--test_split=',
        {outputUri: test_split},
        '", "--test_split_for_batch_prediction_component=',
        {outputPath: test_split_json},
        '", "--downsampled_test_split_for_batch_prediction_component=',
        {outputPath: downsampled_test_split_json},
        '", "--instance_baseline_path=',
        {outputUri: instance_baseline},
        '", "--lro_job_info=',
        {inputValue: root_dir},
        '/{{$.pipeline_job_uuid}}/lro", "--gcp_resources_path=',
        {outputPath: gcp_resources},
        '", "--parse_json=true", "--generate_additional_downsample_test_split=true", "--executor_input={{$.json_escape[1]}}"]}}]}}'
      ]]
