"""
patch to add noisy embeddings per https://arxiv.org/abs/2310.05914
"""

import torch
import transformers.models.mistral.modeling_mistral
from transformers.utils import logging

logger = logging.get_logger(__name__)


def replace_mistral_embeddings_with_uniform_distribution(noise_alpha=5):
    # pylint: disable=duplicate-code
    def noised_embed(orig_embed, noise_alpha, model):
        def new_func(input_ids):
            # during training, we add noise to the embedding
            # during generation, we don't add noise to the embedding
            if model.training:
                embed_init = orig_embed(input_ids)
                dims = torch.tensor(embed_init.size(1) * embed_init.size(2))
                mag_norm = noise_alpha / torch.sqrt(dims)
                return embed_init + torch.zeros_like(embed_init).uniform_(
                    -mag_norm, mag_norm
                )
            return orig_embed(input_ids)

        return new_func

    def post_init(orig_post_init):
        def new_func(self):
            orig_post_init(self)
            self.embed_tokens.forward = noised_embed(
                self.embed_tokens.forward, noise_alpha, self
            )

        return new_func

    transformers.models.mistral.modeling_mistral.MistralModel.post_init = post_init(
        transformers.models.mistral.modeling_mistral.MistralModel.post_init
    )