| 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 | |