name: train_tfhub_model
description: |
  Launch Vertex CustomJob to train a new Cloud natural language TFHub model.

  Args:
      project (str):
          Required. GCP project to run CustomJob.
      location (str):
          Required. Location of the job. Defaults to `us-central1`.
      input_data_path (Optional[str]):
          Required. GCS path to the file where data is stored. Should contain train and validation splits.

          For details on how to format the data, see
          https://cloud.google.com/vertex-ai/docs/text-data/classification/prepare-data.
      input_format (Optional[str]):
          Required. Input data format; supports "csv" and "jsonl". Defaults to "jsonl".

          For details on how to format the data, see
          https://cloud.google.com/vertex-ai/docs/text-data/classification/prepare-data.
      machine_type (Optional[str]):
          Type of machine for running the model training on dedicated resources. Defaults to `n1-highmem-8.`

          For more details about the machine spec, see
          https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec
      accelerator_type (Optional[str]):
          Type of accelerator. Defaults to `NVIDIA_TESLA_T4`.

          For more details about the machine spec, see
          https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec.
      accelerator_count (Optional[int]):
          Number of accelerators. Defaults to `1`.

          For more details about the machine spec, see
          https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec.
      natural_language_task_type (Optional[str]):
          Task type to train the model on.
          Possible values are ["CLASSIFICATION", "MULTILABEL_CLASSIFICATION", "BINARY_CLASSIFICATION"]. Defaults to `CLASSIFICATION`.
      model_architecture (Optional[str]):
          Model you wish to train. Possible values are ["MULTILINGUAL_BERT_BASE_CASED", "ENGLISH_BERT_BASE_UNCASED"]. Defaults to `MULTILINGUAL_BERT_BASE_CASED`.
      train_batch_size (Optional[int]):
          Batch size to use during training.
          Note that increasing this when running on GPUs could cause training failures due to insufficient memory. Defaults to `32`.
      eval_batch_size (Optional[int]):
          Batch size to use during eval.
          Note that increasing this when running on GPUs could cause training failures due to insufficient memory. Defaults to `32`.
      num_epochs (Optional[int]):
          Max number of epochs to train for.
          Early stopping may prevent this number from being reached, but the learning rate schedule is influence by this value regardless. Defaults to `30`.
      dropout (Optional[float]):
          Dropout rate for dropout layer(s). Defaults to `0.35`.
      initial_learning_rate (Optional[float]):
          Starting learning rate for the model. Defaults to `1e-4`.
      warmup (Optional[float]):
          Fraction of training steps that we should slowly raise the learning rate for (i.e. first 10 percent of training steps).
          Note that the "lamb" optimizer will ignore this value as it has built-in warmup. Defaults to `0.1`.
      optimizer_type (Optional[str]):
          Type of optimizer to use. Options are ["lamb", "adamw"]. Defaults to `lamb`.
      random_seed (Optional[int]):
          Random seed to use during training; maintain value for reproducibility. Defaults to `0`.
      train_steps_per_epoch (Optional[int]):
          Training steps per epoch during diintibuted training.
          Only used when there are multiple GPUs. Defaults to `-1`.
      validation_steps_per_epoch (Optional[int]):
          Validation steps per epoch during distributed training.
          Only used when there are multiple GPUs. Defaults to `-1`.
      test_steps_per_epoch (Optional[str]):
          Test steps per epoch during distributed training.
          Only used when there are multiple GPUs. Defaults to `-1`.
  Returns:
      gcp_resources (str):
          Serialized gcp_resources proto tracking the training job.
          For more details on GCP resources proto, see
          https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/proto/README.md.
      model_output (system.Artifact):
          Artifact tracking the batch prediction job output. This is only available if
          gcs_destination_output_uri_prefix is specified.

inputs:
- {name: project, type: String}
- {name: location, type: String}
- {name: input_data_path, type: String}
- {name: input_format, type: String, default: 'JSONL'}
- {name: machine_type, type: String, default: 'n1-highmem-8'}
- {name: accelerator_type, type: String, default: 'NVIDIA_TESLA_T4'}
- {name: accelerator_count, type: Integer, default: '1'}
- {name: natural_language_task_type, type: String, default: 'CLASSIFICATION'}
- {name: model_architecture, type: String, default: 'MULTILINGUAL_BERT_BASE_CASED'}
- {name: train_batch_size, type: Integer, default: '32'}
- {name: eval_batch_size, type: Integer, default: '32'}
- {name: num_epochs, type: Integer, default: '30'}
- {name: dropout, type: Float, default: '0.35'}
- {name: initial_learning_rate, type: Float, default: '1e-4'}
- {name: warmup, type: Float, default: '0.1'}
- {name: optimizer_type, type: String, default: 'lamb'}
- {name: random_seed, type: Integer, default: '0'}
- {name: train_steps_per_epoch, type: Integer, default: '-1'}
- {name: validation_steps_per_epoch, type: Integer, default: '-1'}
- {name: test_steps_per_epoch, type: Integer, default: '-1'}
outputs:
- {name: gcp_resources, type: String}
- {name: model_output, type: TensorflowSavedModel}

implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:latest
    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": "train-tfhub-model-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}",
          "job_spec": {"worker_pool_specs": [{"replica_count":"',
          "1",
          '", "machine_spec": {',
          '"machine_type": "', {inputValue: machine_type}, '"',
          ', "accelerator_type": "', {inputValue: accelerator_type}, '"',
          ', "accelerator_count": ', {inputValue: accelerator_count},
          '}',
          ', "container_spec": {"image_uri":"',
          "us-docker.pkg.dev/vertex-ai-restricted/automl-language/tfhub_model_trainer_image_gpu:latest",
          '", "args": [
              "--input_data_path=', {inputValue: input_data_path}, '",
              "--input_format=', {inputValue: input_format}, '",
              "--natural_language_task_type=', {inputValue: natural_language_task_type}, '",
              "--model_architecture=', {inputValue: model_architecture}, '",
              "--train_batch_size=', {inputValue: train_batch_size}, '",
              "--eval_batch_size=', {inputValue: eval_batch_size}, '",
              "--num_epochs=', {inputValue: num_epochs}, '",
              "--dropout=', {inputValue: dropout}, '",
              "--initial_learning_rate=', {inputValue: initial_learning_rate}, '",
              "--warmup=', {inputValue: warmup}, '",
              "--optimizer_type=', {inputValue: optimizer_type}, '",
              "--random_seed=', {inputValue: random_seed}, '",
              "--train_steps_per_epoch=', {inputValue: train_steps_per_epoch}, '",
              "--validation_steps_per_epoch=', {inputValue: validation_steps_per_epoch}, '",
              "--test_steps_per_epoch=', {inputValue: test_steps_per_epoch}, '",
              "--model_output_path=', {outputPath: model_output}, '"
              ]}}]}}',
        ],
      ]