# 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: automl_tabular_tabnet_trainer
description: |
  Launch a TabNet custom training job using Vertex CustomJob API.

    Args:
        project (str):
            Required. The GCP project that runs the pipeline components.
        location (str):
            Required. The GCP region that runs the pipeline components.
        root_dir (str): Required. The root GCS directory for the pipeline components.
        target_column (str): Required. The target column name.
        prediction_type (str): Required. The type of prediction the model is to produce.
            "classification" or "regression".
        weight_column (Optional[str]): The weight column name.
        max_steps (Optional[int]): Number of steps to run the trainer for.
        max_train_secs (Optional[int]): Amount of time in seconds to run the trainer for.
        learning_rate (float): The learning rate used by the linear optimizer.
        large_category_dim (Optional[int]): Embedding dimension for categorical feature with
            large number of categories.
        large_category_thresh (Optional[int]): Threshold for number of categories to apply large_category_dim
            embedding dimension to.
        yeo_johnson_transform (Optional[bool]): Enables trainable Yeo-Johnson power transform.
        feature_dim (Optional[int]): Dimensionality of the hidden representation in feature
            transformation block.
        feature_dim_ratio (Optional[float]): The ratio of output dimension (dimensionality of the
            outputs of each decision step) to feature dimension.
        num_decision_steps (Optional[int]): Number of sequential decision steps.
        relaxation_factor (Optional[float]): Relaxation factor that promotes the reuse of each feature
            at different decision steps. When it is 1, a feature is enforced to be
            used only at one decision step and as it increases, more flexibility is provided to use a feature at multiple decision steps.
        decay_every (Optional[float]): Number of iterations for periodically applying learning rate decaying.
        gradient_thresh (Optional[float]): Threshold for the norm of gradients for clipping.
        sparsity_loss_weight (Optional[float]): Weight of the loss for sparsity regularization
            (increasing it will yield more sparse feature selection).
        batch_momentum (Optional[float]): Momentum in ghost batch normalization.
        batch_size_ratio (Optional[float]): The ratio of virtual batch size
            (size of the ghost batch normalization) to batch size.
        num_transformer_layers (Optional[int]): The number of transformer layers for each decision step.
            used only at one decision step and as it increases, more flexibility is provided to use a feature at multiple decision steps.
        num_transformer_layers_ratio (Optional[float]): The ratio of shared transformer layer to transformer layers.
        class_weight (Optional[float]): The class weight is used to computes a weighted cross
            entropy which is helpful in classify imbalanced dataset. Only used for classification.
        loss_function_type (Optional[str]): Loss function type. Loss function in classification
            [cross_entropy, weighted_cross_entropy, focal_loss], default is cross_entropy.
            Loss function in regression: [rmse, mae, mse], default is mse.
        alpha_focal_loss (Optional[float]): Alpha value (balancing factor) in focal_loss function. Only used for classification.
        gamma_focal_loss (Optional[float]): Gamma value (modulating factor) for focal loss for focal loss. Only used for classification.
        enable_profiler (Optional[bool]): Enables profiling and saves a trace during evaluation.
        seed (Optional[int]): Seed to be used for this run.
        eval_steps (Optional[int]): Number of steps to run evaluation for. If not
            specified or negative, it means run evaluation on the whole validation
            dataset. If set to 0, it means run evaluation for a fixed number of
            samples.
        batch_size (Optional[int]): Batch size for training.
        eval_frequency_secs (Optional[int]): Frequency at which evaluation and checkpointing will
            take place.
        replica_count (Optional[int]): The replica count.
        machine_spec (Optional[Dict[str, Any]]): The machine spec.
            See https://cloud.google.com/compute/docs/machine-types for options.
        instance_baseline (AutoMLTabularInstanceBaseline): The path to a JSON file for baseline values.
        metadata (TabularExampleGenMetadata): Amount of time in seconds to run the trainer for.
        materialized_train_split (MaterializedSplit): The path to the materialized train split.
        materialized_eval_split (MaterializedSplit): The path to the materialized validation split.
        materialized_test_split (MaterializedSplit): The path to the materialized test split.
        transform_output (TransformOutput): The path to transform output.
        training_schema_uri (TrainingSchema): The path to the training schema.
        encryption_spec_key_name (Optional[str]): The KMS key name.

    Returns:
        gcp_resources (str):
            Serialized gcp_resources proto tracking the custom training job.
        unmanaged_container_model (google.UnmanagedContainerModel): The UnmanagedContainerModel artifact.

inputs:
  - { name: project, type: String }
  - { name: location, type: String }
  - { name: root_dir, type: String }
  - { name: target_column, type: String }
  - { name: prediction_type, type: String }
  - { name: weight_column, type: String, default: "" }
  - { name: max_steps, type: Integer, default: -1 }
  - { name: max_train_secs, type: Integer, default: -1 }
  - { name: learning_rate, type: Float }
  - { name: large_category_dim, type: Integer, default: 1 }
  - { name: large_category_thresh, type: Integer, default: 300 }
  - { name: yeo_johnson_transform, type: Boolean, default: "true" }
  - { name: feature_dim, type: Integer, default: 64 }
  - { name: feature_dim_ratio, type: Float, default: 0.5 }
  - { name: num_decision_steps, type: Integer, default: 6 }
  - { name: relaxation_factor, type: Float, default: 1.5 }
  - { name: decay_every, type: Float, default: 100}
  - { name: gradient_thresh, type: Float, default: 2000 }
  - { name: sparsity_loss_weight, type: Float, default: 0.00001 }
  - { name: batch_momentum, type: Float, default: 0.95 }
  - { name: batch_size_ratio, type: Float, default: 0.25 }
  - { name: num_transformer_layers, type: Integer, default: 4 }
  - { name: num_transformer_layers_ratio, type: Float, default: 0.25 }
  - { name: class_weight, type: Float, default: 1.0 }
  - { name: loss_function_type, type: String, default: 'default' }
  - { name: alpha_focal_loss, type: Float, default: 0.25 }
  - { name: gamma_focal_loss, type: Float, default: 2.0 }
  - { name: enable_profiler, type: Boolean, default: "false" }
  - { name: seed, type: Integer, default: 1 }
  - { name: eval_steps, type: Integer, default: 0 }
  - { name: batch_size, type: Integer, default: 100 }
  - { name: eval_frequency_secs, type: Integer, default: 600 }
  - { name: replica_count, type: Integer, default: 1 }
  - { name: machine_spec, type: JsonObject, default: '{"machine_type": "c2-standard-16"}' }
  - { name: instance_baseline, type: AutoMLTabularInstanceBaseline }
  - { name: metadata, type: TabularExampleGenMetadata }
  - { name: materialized_train_split, type: MaterializedSplit }
  - { name: materialized_eval_split, type: MaterializedSplit }
  - { name: materialized_test_split, type: MaterializedSplit, default: "" }
  - { name: transform_output, type: TransformOutput }
  - { name: training_schema_uri, type: TrainingSchema }
  - { name: encryption_spec_key_name, type: String, default: "" }

outputs:
  - { name: gcp_resources, type: String }
  - { name: unmanaged_container_model, type: google.UnmanagedContainerModel }

implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:1.10.13
    command: [python3, -u, -m, google_cloud_pipeline_components.container.v1.gcp_launcher.launcher]
    args: [
      --type, CustomJob,
      --project, { inputValue: project },
      --location, { inputValue: location },
      --gcp_resources, { outputPath: gcp_resources },
      --payload,
      concat: [
          '{"display_name": "automl-tabular-tabnet-trainer-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}", "encryption_spec": {"kms_key_name":"',
          { inputValue: encryption_spec_key_name },
          '"}, "job_spec": {"worker_pool_specs": [{"replica_count":"',
          { inputValue: replica_count },
          '", "machine_spec": ',
          { inputValue: machine_spec },
          ', "container_spec": {"image_uri":"',
          'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/tabnet-training:20220629_0725_RC00',
          '", "args": ["--target_column=',
          { inputValue: target_column },
          '", "--weight_column=',
          { inputValue: weight_column },
          '", "--model_type=',
          { inputValue: prediction_type },
          '", "--prediction_docker_uri=',
          'us-docker.pkg.dev/vertex-ai/automl-tabular/prediction-server:20220629_0725_RC00',
          '", "--baseline_path=',
          { inputUri: instance_baseline },
          '", "--metadata_path=',
          { inputUri: metadata },
          '", "--transform_output_path=',
          { inputUri: transform_output },
          '", "--training_schema_path=',
          { inputUri: training_schema_uri },
          '", "--job_dir=',
          { inputValue: root_dir },
          '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/train", "--training_data_path=',
          { inputUri: materialized_train_split },
          '", "--validation_data_path=',
          { inputUri: materialized_eval_split },
          '", "--test_data_path=',
          { inputUri: materialized_test_split },
          '", "--max_steps=',
          { inputValue: max_steps },
          '", "--max_train_secs=',
          { inputValue: max_train_secs },
          '", "--learning_rate=',
          { inputValue: learning_rate },
          '", "--large_category_dim=',
          { inputValue: large_category_dim },
          '", "--large_category_thresh=',
          { inputValue: large_category_thresh },
          '", "--yeo_johnson_transform=',
          { inputValue: yeo_johnson_transform },
          '", "--feature_dim=',
          { inputValue: feature_dim },
          '", "--feature_dim_ratio=',
          { inputValue: feature_dim_ratio },
          '", "--num_decision_steps=',
          { inputValue: num_decision_steps },
          '", "--relaxation_factor=',
          { inputValue: relaxation_factor },
          '", "--decay_every=',
          { inputValue: decay_every },
          '", "--gradient_thresh=',
          { inputValue: gradient_thresh },
          '", "--sparsity_loss_weight=',
          { inputValue: sparsity_loss_weight },
          '", "--batch_momentum=',
          { inputValue: batch_momentum },
           '", "--batch_size_ratio=',
          { inputValue: batch_size_ratio },
          '", "--num_transformer_layers=',
          { inputValue: num_transformer_layers },
          '", "--num_transformer_layers_ratio=',
          { inputValue: num_transformer_layers_ratio },
          '", "--class_weight=',
          { inputValue: class_weight },
          '", "--loss_function_type=',
          { inputValue: loss_function_type },
          '", "--alpha_focal_loss=',
          { inputValue: alpha_focal_loss },
          '", "--gamma_focal_loss=',
          { inputValue: gamma_focal_loss },
          '", "--enable_profiler=',
          { inputValue: enable_profiler },
          '", "--seed=',
          { inputValue: seed },
          '", "--eval_steps=',
          { inputValue: eval_steps },
          '", "--batch_size=',
          { inputValue: batch_size },
          '", "--eval_frequency_secs=',
          { inputValue: eval_frequency_secs },
          '", "--executor_input={{$.json_escape[1]}}"]}}]}}',
      ]]
