description: GPT2 COMET 2020 Pipeline - COMET TransOMCS
tasks:
  - name: gpt2-comet-transomcs
    cluster: ai2/on-prem-ai2-server2
    spec:
      env:
        TRAIN_BATCH_SIZE: 32  # input batch size for training (default: 64)
        VALID_BATCH_SIZE: 32   # input batch size for testing (default: 1000)
        TRAIN_EPOCHS: 1        # number of epochs to train (default: 10)
        VAL_EPOCHS: 1 
        LEARNING_RATE: 0.00005   # learning rate (default: 0.01)
        SEED: 42               # random seed (default: 42)
        TRAIN_DATA_PATH: "/data/transomcs_train.tsv"
        DEV_DATA_PATH: "/data/transomcs_dev.tsv"
        TEST_DATA_PATH: "/data/transomcs_test.tsv"
        IN_LEN: 8
        OUT_LEN: 16
        GPT2_MODEL: "/gpt2-xl"
      image: {{ .Env.IMAGE }}
      resultPath: /models
      requirements:
        gpuCount: 1
      datasetMounts:
        - datasetId: ds_yi6ma1pwef76
          containerPath: /data
        - datasetId: ds_87rekfgmu7fp
          containerPath: /gpt2-xl
      args:
        - python
        - comet_gpt2.py