def get_lr(self):
        """Learning rate decay functions from:
              https://openreview.net/pdf?id=BJYwwY9ll pg. 4"""

        # Use linear warmup for the initial part.
        if self.warmup_steps > 0 and self.num_steps <= self.warmup_steps:
            if self.num_steps == self.warmup_steps:
                self.warmup_tokens = self.consumed_tokens
            return self.max_lr * float(self.num_steps) / \
                float(self.warmup_steps)

        # If the learning rate is constant, just return the initial value.
        if self.decay_style == 'constant':
            return self.max_lr

        if self.consumed_tokens > self.decay_tokens:
            return self.min_lr
        consumed_tokens_ = self.consumed_tokens - self.warmup_tokens
        decay_tokens_ = self.decay_tokens - self.warmup_tokens
        decay_ratio = float(consumed_tokens_) / float(decay_tokens_)
        assert decay_ratio >= 0.0
        assert decay_ratio <= 1.0
        delta_lr = self.max_lr - self.min_lr

        if self.decay_style == 'linear':
            coeff = (1.0 - decay_ratio)
        elif self.decay_style == 'cosine':
            coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)
        else:
            raise Exception('{} decay style is not supported.'.format(
                self.decay_style))

        return self.min_lr + coeff * delta_lr
Learning rate decay functions from:
              https://openreview.net/pdf?id=BJYwwY9ll pg. 4
training/data_efficiency/gpt_finetuning/learning_rates.py