---
trainer:
  class_path: eva.Trainer
  init_args:
    n_runs: &N_RUNS ${oc.env:N_RUNS, 20}
    default_root_dir: &OUTPUT_ROOT ${oc.env:OUTPUT_ROOT, logs/${oc.env:MODEL_NAME, vitg14_reg}/offline/camelyon16}
    max_epochs: &MAX_EPOCHS ${oc.env:MAX_EPOCHS, 100}
    checkpoint_type: ${oc.env:CHECKPOINT_TYPE, best}
    accelerator: ${oc.env:ACCELERATOR, auto}
    devices: ${oc.env:NUM_DEVICES, 1}
    callbacks:
      - class_path: eva.callbacks.ConfigurationLogger
      - class_path: lightning.pytorch.callbacks.TQDMProgressBar
        init_args:
          refresh_rate: ${oc.env:TQDM_REFRESH_RATE, 1}
      - class_path: lightning.pytorch.callbacks.LearningRateMonitor
        init_args:
          logging_interval: epoch
      - class_path: lightning.pytorch.callbacks.ModelCheckpoint
        init_args:
          filename: best
          save_last: ${oc.env:SAVE_LAST, false}
          save_top_k: 1
          monitor: &MONITOR_METRIC ${oc.env:MONITOR_METRIC, val/BinaryBalancedAccuracy}
          mode: &MONITOR_METRIC_MODE ${oc.env:MONITOR_METRIC_MODE, max}
      - class_path: lightning.pytorch.callbacks.EarlyStopping
        init_args:
          min_delta: 0
          patience: ${oc.env:PATIENCE, 20}
          monitor: *MONITOR_METRIC
          mode: *MONITOR_METRIC_MODE
      - class_path: eva.callbacks.ClassificationMultiEmbeddingsInMemoryWriter
        init_args:
          output_dir: &DATASET_EMBEDDINGS_ROOT ${oc.env:EMBEDDINGS_ROOT, ./data/embeddings/${oc.env:MODEL_NAME, vitg14_reg}/camelyon16}
          save_every_n: 10_000
          dataloader_idx_map:
            0: train
            1: val
            2: test
          metadata_keys: ["wsi_id"]  
          backbone:
            class_path: eva.core.models.wrappers.ModelFromFunction
            init_args:
              path: torch.hub.load
              arguments:
                repo_or_dir: facebookresearch/dinov2 
                model: dinov2_vitg14_reg
                pretrained: true
              checkpoint_path: ${oc.env:CHECKPOINT_PATH, ../checkpoints/teacher_epoch250000.pth}

          overwrite: true
    logger:
      - class_path: lightning.pytorch.loggers.TensorBoardLogger
        init_args:
          save_dir: *OUTPUT_ROOT
          name: ""
model:
  class_path: eva.HeadModule
  init_args:
    head:
      class_path: eva.vision.models.networks.ABMIL
      init_args:
        input_size: ${oc.env:IN_FEATURES, 1536}
        output_size: &NUM_CLASSES 1
        projected_input_size: 128
    criterion: torch.nn.BCEWithLogitsLoss
    optimizer:
      class_path: torch.optim.AdamW
      init_args:
        lr: ${oc.env:LR_VALUE, 0.001}
        betas: [0.9, 0.999]
    metrics:
      common:
        - class_path: eva.metrics.AverageLoss
        - class_path: eva.metrics.BinaryClassificationMetrics
data:
  class_path: eva.DataModule
  init_args:
    datasets:
      train:
        class_path: eva.datasets.InMemoryMultiEmbeddingsClassificationDataset
        init_args: &DATASET_ARGS
          root: *DATASET_EMBEDDINGS_ROOT
          manifest_file: manifest.csv
          split: train
          embeddings_transforms:
            class_path: eva.core.data.transforms.Pad2DTensor
            init_args:
              pad_size: &N_PATCHES ${oc.env:N_PATCHES, 1000}
          target_transforms:
            class_path: torchvision.transforms.v2.ToDtype
            init_args:
              dtype: torch.float32
      val:
        class_path: eva.datasets.InMemoryMultiEmbeddingsClassificationDataset
        init_args:
          <<: *DATASET_ARGS
          split: val
      test:
        class_path: eva.datasets.InMemoryMultiEmbeddingsClassificationDataset
        init_args:
          <<: *DATASET_ARGS
          split: test
      predict:
        - class_path: eva.vision.datasets.Camelyon16
          init_args: &PREDICT_DATASET_ARGS
            root: ${oc.env:DATA_ROOT, /block/eva-data/camelyon16}
            sampler:
              class_path: eva.vision.data.wsi.patching.samplers.ForegroundGridSampler
              init_args:
                max_samples: *N_PATCHES
            width: 224
            height: 224
            target_mpp: 0.25
            split: train
            coords_path: ${data.init_args.datasets.train.init_args.root}/coords_${.split}.csv
            image_transforms:
              class_path: eva.vision.data.transforms.common.ResizeAndCrop
              init_args:
                size: ${oc.env:RESIZE_DIM, 224} 
                mean: ${oc.env:NORMALIZE_MEAN, [0.485, 0.456, 0.406]} 
                std: ${oc.env:NORMALIZE_STD, [0.229, 0.224, 0.225]}
        - class_path: eva.vision.datasets.Camelyon16
          init_args:
            <<: *PREDICT_DATASET_ARGS
            split: val
        - class_path: eva.vision.datasets.Camelyon16
          init_args:
            <<: *PREDICT_DATASET_ARGS
            split: test
    dataloaders:
      train:
        batch_size: &BATCH_SIZE ${oc.env:BATCH_SIZE, 32}
        num_workers: &N_DATA_WORKERS ${oc.env:N_DATA_WORKERS, 4}
        shuffle: true
      val:
        batch_size: *BATCH_SIZE
        num_workers: *N_DATA_WORKERS
      test:
        batch_size: *BATCH_SIZE
        num_workers: *N_DATA_WORKERS
      predict:
        batch_size: &PREDICT_BATCH_SIZE ${oc.env:PREDICT_BATCH_SIZE, 64}
        num_workers: *N_DATA_WORKERS