##################
# Trainer settings
##################


TASK: UniCLTask

NAME: 'Example Eval Configuration'
SAVE_TIMER_LOG: true

# TUTORIAL STEP 1: CHOOSE SAVE DIR
SAVE_DIR: ''
LOG_EVERY: 10
LOGLEVEL_OVERRIDE: INFO
LOG_GPU_MEM: true
RESUME: False
RESET_DATA_LOADER: false

FP16: true
ZERO_STAGE: 0
DEEPSPEED: false
# ZERO_STAGE: 1
AMP: PYTORCH
# USE_APEX_DDP: false
# USE_APEX_AMP: false
# USE_HIT: false

FIND_UNUSED_PARAMETERS: false

SAVE_PER_OPTIM_STEPS: 500
EVAL_PER_OPTIM_STEPS: 250
EVAL_AT_START: False
# SAVE_PER_UPDATE_NUM: -1
# EVAL_PER_UPDATE_NUM: 0 # 0: do evaluation when saving checkpoint, -1: don't do evaluation

NO_AUTO_LR_SCALING: true
GRAD_CLIPPING: 1.0 #0.07

SET_SAMPLER_EPOCH: true

DONT_LOAD_MODEL: true

user_dir: "./MainzVision" # lower case due to it is used in mainz as such

##################
# Task settings
##################



VERBOSE: true
WORKERS: 6
PIN_MEMORY: true
IMAGE_ENCODER:
  NAME: davit_v1
  NUM_CLASSES: 0
  #IMAGE_SIZE: [384, 384]
  IMAGE_SIZE: [480, 480]
  LOAD_PRETRAINED: true
  PRETRAINED: ''
  PRETRAINED_LAYERS: '*'
  IMAGE_MEAN: [0.485, 0.456, 0.406]
  IMAGE_STD: [0.229, 0.224, 0.225]
  SPEC:
    DROP_RATE: 0.1
    DROP_PATH_RATE: 0.2
    PATCH_SIZE: [7, 3, 3, 3]
    PATCH_STRIDE: [4, 2, 2, 2]
    PATCH_PADDING: [3, 1, 1, 1]
    PATCH_PRENORM: [false, true, true, true]
    DIM_EMBED: [256, 512, 1024, 2048]
    NUM_HEADS: [8, 16, 32, 64]
    NUM_GROUPS: [8, 16, 32, 64]
    DEPTHS: [1, 1, 9, 1]
    WINDOW_SIZE: 12
    ENABLE_CHECKPOINT: true

LANG_ENCODER:
  NAME: transformer
  LOAD_PRETRAINED: false
  PRETRAINED: ''
  PRETRAINED_LAYERS: '*'
  TOKENIZER: clip
  CONTEXT_LENGTH: 77
  WIDTH: 1024
  HEADS: 16
  LAYERS: 16
  AUTOGRESSIVE: false

UNICL_MODEL:
  DIM_PROJECTION: 1024
  GATHER_TENSORS: true
  LOAD_PRETRAINED: true

  # TUTORIAL STEP 2: CHOOSE MODEL PATH
  PRETRAINED: ''

  PRETRAINED_LAYERS: '*'

AUG:
  MIXUP_PROB: 0.0
  MIXUP: 0.8
  MIXCUT: 1.0
  MIXCUT_MINMAX: []
  MIXUP_SWITCH_PROB: 0.5
  MIXUP_MODE: 'batch'
  SCALE: [0.8, 1.0]
  RATIO: [0.75, 1.3333333]
  INTERPOLATION: 'bicubic'
  TORCHVISION_AUG:
    AUTO_AUGMENT: ta_wide
    RE_PROB: 0.25
    HFLIP: 0.0
    VFLIP: 0.0

LOSS:
  LOSS: UniCL
DATASET:
  DATASET: 'image_text_pairs_v2'
  TEXT_FORMAT: 'json'
  ROOT: ''
  TRAIN_SET: 'mimic_cxr_v2-chestxray14-chexpertv4-irma2009_v2-rsnaboneage-mura-bingmedicalfewshot'
  DATA_FORMAT: 'tsv'
  SAMPLER: 'default'
  LOADER: 'default'
  TOKEN_FILE: ''
  #PROMPT_ENGINEERING: False
  #SAMPLER: 'chunk'
  #LOADER: 'azcopy'
  #TOKEN_FILE: 'cliptrainingpairs.txt'
  #TEST_SET: 'MarsAtrain'


# TUTORIAL STEP 3: CHOOSE ALL BELOW EVAL PATHS (THESE ARE ALL OPTIONAL EXTRA EVALS)
# Note how one eval is ZIP format and the other is TSV format.




EVALDATASET_LTCXR_S100_N100_TEXT_CLASSIFIER:
  TEXT_FORMAT: json
  FORMAT: 'zip'
  SPLIT: 'NIH-CXR-LT'
  ZIP_FILE: ''
  ZIP_MAP_FILE: ''
  LABEL_FILE: ''
  IMAGE_TSV: ''
  TEXT_TSV: ''
  CWEIGHT_FILE: ''
  ZS_MODE: 2
  ZS_WEIGHT: 1.0
  KNN: 100
#  CLASSIFICATION_SETS: ['NIH-CXR-LT']
#  NUM_CLASSES: [20]




# TUTORIAL STEP 4: SET THE DEFAULT ZEROSHOT EVAL (THIS IS THE MANDATORY EVAL)

ZEROSHOT_EVAL_DATASET:
  FORMAT: 'zip'
  SPLIT: 'NIH-CXR-LT'
  ZIP_FILE: ''
  ZIP_MAP_FILE: ''
  LABEL_FILE: ''



EVALUATION_SPLITS: ['cls-zeroshot-eval']
TEST:
  BATCH_SIZE_PER_GPU: 8
  MODEL_FILE: ''
  CENTER_CROP: false
TRAIN:
  BATCH_SIZE_TOTAL: 1024
  BATCH_SIZE_PER_GPU: 16

  SHUFFLE: true

WEIGHT_SMOOTHING:
  decay: 0.999
  use_cpu: False
  eval_smoothed_weight: True

START_LEARNING_RATE: 0.00001
# MAX_NUM_EPOCHS: 2
MAX_NUM_EPOCHS: 100
OPTIMIZER: AdamW # adam
OPTIMIZER_PARAMS:
  weight_decay: 0.2 #0.1
CUSTOMIZED_PARAMS_CONF:
  NO_WEIGHT_DECAY_MODULES: ['dw', 'norm']
  WEIGHT_DECAY_PATTERNS:
    "\\.bias$": 0.0
    "logit_scale": 0.0
    "positional_embedding": 0.0
    "token_embedding": 0.0



LR_SCHEDULER: TimmScheduler
LR_SCHEDULER_PARAMS:
  sched: cosine
  warmup_steps: 5
  warmup_lr: 0.000000001
  min_lr: 0.000000001

# GRADIENT_ACCUMULATE_STEP will be updated by:
# BATCH_SIZE_TOTAL // (BATCH_SIZE_PER_GPU * world_size)
GRADIENT_ACCUMULATE_STEP: -1