# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/utils/02_ML_Tools.ipynb.

# %% auto 0
__all__ = ['LRSchedulerPerStep']

# %% ../nbs/utils/02_ML_Tools.ipynb 3
# General libraries:
# import numpy as np

import wandb
import datetime
import torch


# %% ../nbs/utils/02_ML_Tools.ipynb 4
class LRSchedulerPerStep():
    """
    Learning rate scheduler from Attention is all you need paper
    Change: added basic LR as parameter rather than automated based on model dimension

    """
    
    def __init__(self, optimizer, base_learning_rate=0.0001, warmup=4000):

        # Ensure optimizer is a PyTorch optimizer
        if not isinstance(optimizer, torch.optim.Optimizer):
            raise ValueError('Optimizer must be a PyTorch optimizer')
        
        self.optimizer = optimizer
        self.basic = base_learning_rate
        self.warm = warmup**-1.5
        self.step_num = 0
        self.scaling_factor = 1/warmup**-0.5 # ensures that the peak realtive to the base lr is always 1

        print("set up schedule")
        
    def step(self):
        self.step_num += 1
        lr = self.basic * self.scaling_factor * min(self.step_num**-0.5, self.step_num*self.warm)
        
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr
		
