#!/usr/bin/env bash

# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# 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.
# ==============================================================================

# CNN-DM Summarization configurations for NVIDIA DGX A100 (8x NVIDIA A100 40GB GPU)

dgxa100_8gpu_fp16 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  LR=1.25e-4
  BS=40
  ACCUM=1
  PRECISION="fp16"
  TRAIN_STEPS=2000
  WARMUP_STEPS=50
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=128
  PRED_BS=128
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgxa100_8gpu_fp16_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  PRECISION="fp16"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=128

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}

dgxa100_8gpu_tf32 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  LR=1.25e-4
  BS=24
  ACCUM=1
  PRECISION="tf32"
  TRAIN_STEPS=3333
  WARMUP_STEPS=50
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=128
  PRED_BS=64
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgxa100_8gpu_tf32_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  PRECISION="tf32"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=64

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}


# CNN-DM Summarization configurations for NVIDIA DGX-2H (16x NVIDIA V100 32GB GPU)

dgx2_16gpu_fp16 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=16
  LR=1e-4
  BS=14
  ACCUM=1
  PRECISION="fp16"
  TRAIN_STEPS=3500
  WARMUP_STEPS=100
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=128
  PRED_BS=96
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgx2_16gpu_fp16_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=16
  PRECISION="fp16"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=96

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}

dgx2_16gpu_fp32 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=16
  LR=7e-5
  BS=8
  ACCUM=1
  PRECISION="fp32"
  TRAIN_STEPS=6125
  WARMUP_STEPS=150
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=96
  PRED_BS=32
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgx2_16gpu_fp32_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=16
  PRECISION="fp32"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=32

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}

# CNN-DM Summarization configurations for NVIDIA DGX-1 (8x NVIDIA V100 32GB GPU)

dgx1_8gpu_fp16 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  LR=5.5e-5
  BS=14
  ACCUM=1
  PRECISION="fp16"
  TRAIN_STEPS=7000
  WARMUP_STEPS=150
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=128
  PRED_BS=96
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgx1_8gpu_fp16_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  PRECISION="fp16"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=96

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}

dgx1_8gpu_fp32 ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  LR=5.5e-5
  BS=8
  ACCUM=1
  PRECISION="fp32"
  TRAIN_STEPS=12250
  WARMUP_STEPS=150
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  EVAL_BS=96
  PRED_BS=32
  VAL_CHECK_INTERVAL=0.3

  echo $DATA_DIR $CONFIG_PATH $NUM_GPU $LR $BS $ACCUM $PRECISION $TRAIN_STEPS $WARMUP_STEPS $MAX_SOURCE_LEN $MAX_TARGET_LEN $EVAL_BEAMS $EVAL_BS $PRED_BS $VAL_CHECK_INTERVAL
}

dgx1_8gpu_fp32_eval ()
{
  DATA_DIR=data/cnn_dm/
  CONFIG_PATH="configs/config.json"
  NUM_GPU=8
  PRECISION="fp32"
  MAX_SOURCE_LEN=1024
  MAX_TARGET_LEN=142
  EVAL_BEAMS=4
  PRED_BS=32

  echo $PRED_BS $NUM_GPU $PRECISION $EVAL_BEAMS $MAX_SOURCE_LEN $MAX_TARGET_LEN $DATA_DIR $CONFIG_PATH
}
