Method lr_lambda

fireredtts2/llm/utils.py:33–59

(self, step: int)

        super().__init__(optimizer, self.lr_lambda, last_epoch=-1)

    def lr_lambda(self, step: int) -> float:
        if step < self.warmup_steps:
            return step / self.warmup_steps
        else:
            if self.decay_type == "linear":
                return (self.total_steps - step) / (
                    self.total_steps - self.warmup_steps
                )
            elif self.decay_type == "constant":
                return 1.0
            elif self.decay_type == "exponential":
                return 0.1 ** (
                    (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
                )
            elif self.decay_type == "cosine":
                return 0.5 * (
                    1
                    + torch.cos(
                        torch.pi
                        * torch.tensor(
                            (step - self.warmup_steps)
                            / (self.total_steps - self.warmup_steps)
                        )
                    )
                )
            else:
                raise ValueError(f"Invalid decay type: {self.decay_type}")


additional_special_tokens = [