name: train_tensorflow_text_classification_model
description: |
  Creates a trained text classification TensorFlow model.
  Args:
      preprocessed_training_data_path (str):
          Path for the training data for text classification in jsonl format. The data must contain
          'text' (input) and 'label' (output for a prediction) fields.
      preprocessed_validation_data_path (str):
          Path for the validation data (same format as the preprocessed_training_data).
      trained_model_path (str):
          Output path for a trained TensorFlow SavedModel compliant with
          https://github.com/tensorflow/tensorflow/blob/v2.11.0/tensorflow/python/keras/saving/save.py
      class_names (Sequence[str]):
          Sequence of strings of categories for classification.
      model_name (Optional[str]):
          Name of pre-trained BERT encoder model (https://tfhub.dev/google/collections/bert/1) to be used.
          Eligible model_name:
          - bert_en_uncased_L-12_H-768_A-12
          - bert_en_cased_L-12_H-768_A-12
          - bert_multi_cased_L-12_H-768_A-12
          - small_bert/bert_en_uncased_L-2_H-128_A-2 (default)
          - small_bert/bert_en_uncased_L-2_H-256_A-4
          - small_bert/bert_en_uncased_L-2_H-512_A-8
          - small_bert/bert_en_uncased_L-2_H-768_A-12
          - small_bert/bert_en_uncased_L-4_H-128_A-2
          - small_bert/bert_en_uncased_L-4_H-256_A-4
          - small_bert/bert_en_uncased_L-4_H-512_A-8
          - small_bert/bert_en_uncased_L-4_H-768_A-12
          - small_bert/bert_en_uncased_L-6_H-128_A-2
          - small_bert/bert_en_uncased_L-6_H-256_A-4
          - small_bert/bert_en_uncased_L-6_H-512_A-8
          - small_bert/bert_en_uncased_L-6_H-768_A-12
          - small_bert/bert_en_uncased_L-8_H-128_A-2
          - small_bert/bert_en_uncased_L-8_H-256_A-4
          - small_bert/bert_en_uncased_L-8_H-512_A-8
          - small_bert/bert_en_uncased_L-8_H-768_A-12
          - small_bert/bert_en_uncased_L-10_H-128_A-2
          - small_bert/bert_en_uncased_L-10_H-256_A-4
          - small_bert/bert_en_uncased_L-10_H-512_A-8
          - small_bert/bert_en_uncased_L-10_H-768_A-12
          - small_bert/bert_en_uncased_L-12_H-128_A-2
          - small_bert/bert_en_uncased_L-12_H-256_A-4
          - small_bert/bert_en_uncased_L-12_H-512_A-8
          - small_bert/bert_en_uncased_L-12_H-768_A-12
          - albert_en_base
          - electra_small
          - electra_base
          - experts_pubmed
          - experts_wiki_books
          - talking-heads_base
      batch_size (Optional[int]):
          A number of samples processed before the model is updated (must be >= 1 and <= number of
          samples in the dataset).
      num_epochs (Optional[int]):
          Number of training iterations over data.
      learning_rate (Optional[float]): Learning rate controls how quickly the model is adapted to
          the problem, often in the range between 0.0 and 1.0.
      random_seed (Optional[int]):
          The global random seed to ensure the system gets a unique random sequence that is
          deterministic (https://www.tensorflow.org/api_docs/python/tf/random/set_seed).
inputs:
- {name: preprocessed_training_data_path, type: JSONLines, description: Path for the
    training data}
- {name: preprocessed_validation_data_path, type: JSONLines, description: Path for
    the validation data}
- {name: class_names, type: JsonArray, description: List of categories for classification}
- {name: model_name, type: String, description: Name of pre-trained BERT model to
    be used, default: small_bert/bert_en_uncased_L-2_H-128_A-2, optional: true}
- {name: batch_size, type: Integer, description: Batch size, default: '256', optional: true}
- {name: num_epochs, type: Integer, description: Number of epochs, default: '10',
  optional: true}
- {name: learning_rate, type: Float, description: Learning rate, default: '0.0003',
  optional: true}
- {name: random_seed, type: Integer, description: Random seed, default: '0', optional: true}
outputs:
- {name: trained_model_path, type: TensorflowSavedModel, description: Output path
    for trained Tf model}

implementation:
  container:
    image: us-docker.pkg.dev/vertex-ai/ready-to-go-text-classification/training:v0.1
    command: [
      python3,
      /pipelines/component/src/component.py,
      --preprocessed-training-data-path,
      {inputPath: preprocessed_training_data_path},
      --preprocessed-validation-data-path,
      {inputPath: preprocessed_validation_data_path},
      --class-names,
      {inputValue: class_names},
      --model-name,
      {inputValue: model_name},
      --batch-size,
      {inputValue: batch_size},
      --num-epochs,
      {inputValue: num_epochs},
      --learning-rate,
      {inputValue: learning_rate},
      --random-seed,
      {inputValue: random_seed},
      --trained-model-path,
      {outputPath: trained_model_path},
    ]
