config.py

import torch

# Set device cuda for GPU if it is available, otherwise run on the CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Training hyperparameters
batch_size = 50
min_epochs = 1
max_epochs = 10
learning_rate = 5e-5
unfreeze_ratio = 0.2
temperature = 0.1

# Dataset
split_ratio = 0.2
num_workers = 1
num_pos = 20
num_neg = 200
batch_size_simcse = 200

# Compute related
accelerator = "gpu"
devices = 1  # number of gpus
precision = "16-mixed"
log_every_n_steps = 20