# Copyright (c) ModelScope Contributors. All rights reserved.
import torch.nn.functional as F
from torch import Tensor
from transformers import PreTrainedModel
from types import MethodType

from swift.template import TemplateType
from swift.utils import get_logger
from ..constant import LLMModelType
from ..model_arch import ModelArch
from ..model_meta import Model, ModelGroup, ModelMeta
from ..register import ModelLoader, register_model

logger = get_logger()


class BaichuanLoader(ModelLoader):

    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
        model = super().get_model(model_dir, *args, **kwargs)
        # baichuan-13b does not implement the `get_input_embeddings` function
        # fix gradient_checkpointing bug
        try:
            model.get_input_embeddings()
        except NotImplementedError:
            model.__class__.get_input_embeddings = lambda self: self.model.embed_tokens
        return model


register_model(
    ModelMeta(
        LLMModelType.baichuan, [
            ModelGroup([
                Model('baichuan-inc/Baichuan-13B-Chat', 'baichuan-inc/Baichuan-13B-Chat'),
                Model('baichuan-inc/Baichuan-13B-Base', 'baichuan-inc/Baichuan-13B-Base'),
                Model('baichuan-inc/baichuan-7B', 'baichuan-inc/Baichuan-7B'),
            ]),
        ],
        BaichuanLoader,
        template=TemplateType.baichuan,
        architectures=['BaichuanForCausalLM', 'BaiChuanForCausalLM'],
        model_arch=ModelArch.baichuan,
        requires=['transformers<4.34']))


class BaichuanM1Loader(BaichuanLoader):

    def get_model(self, model_dir: str, *args, **kwargs) -> PreTrainedModel:
        from transformers.dynamic_module_utils import get_class_from_dynamic_module
        rotary_embedding = get_class_from_dynamic_module('modeling_baichuan.RotaryEmbedding', model_dir)
        _old_forward = rotary_embedding.forward

        def _new_forward(self, q, k, seqlen_offset=None, cu_seqlens=None, max_seqlen=None):
            q = q.to(k.dtype)
            res = _old_forward(self, q, k, seqlen_offset, cu_seqlens, max_seqlen)
            return res

        rotary_embedding.forward = _new_forward
        return super().get_model(model_dir, *args, **kwargs)


register_model(
    ModelMeta(
        LLMModelType.baichuan_m1, [
            ModelGroup([
                Model('baichuan-inc/Baichuan-M1-14B-Instruct', 'baichuan-inc/Baichuan-M1-14B-Instruct'),
            ]),
        ],
        BaichuanM1Loader,
        template=TemplateType.baichuan_m1,
        architectures=['BaichuanM1ForCausalLM'],
        model_arch=ModelArch.baichuan,
        requires=['transformers>=4.48']))


def patch_baichuan2_lm_head_forward(self, hidden_states: Tensor) -> Tensor:
    # patch: baichuan2 lm_head (fp32 bug)
    if self.training:
        norm_weight = F.normalize(self.weight).to(self.weight.dtype)
    elif self.first_flag:
        self.first_flag = False
        self.weight.data = F.normalize(self.weight).to(self.weight.dtype)
        norm_weight = self.weight
    else:
        norm_weight = self.weight
    return F.linear(hidden_states, norm_weight)


class Baichuan2Loader(ModelLoader):

    def get_model(self, model_dir: str, config, *args, **kwargs) -> PreTrainedModel:
        if not hasattr(config, 'z_loss_weight'):
            config.z_loss_weight = 0
        # patch: baichuan2_13b configuration_baichuan.py bug
        if hasattr(config, 'gradient_checkpointing'):
            gradient_checkpointing = config.gradient_checkpointing
            if isinstance(gradient_checkpointing, (tuple, list)):
                config.gradient_checkpointing = gradient_checkpointing[0]
        model = super().get_model(model_dir, config, *args, **kwargs)
        model_ori = model
        if not hasattr(model, 'lm_head'):  # fix awq
            model = model.model
        new_forward = MethodType(patch_baichuan2_lm_head_forward, model.lm_head)
        if hasattr(model, '_old_forward'):  # device_map
            model.lm_head._old_forward = new_forward
        else:
            model.lm_head.forward = new_forward
        return model_ori


register_model(
    ModelMeta(
        LLMModelType.baichuan2,
        [
            ModelGroup([
                Model('baichuan-inc/Baichuan2-7B-Chat', 'baichuan-inc/Baichuan2-7B-Chat'),
                Model('baichuan-inc/Baichuan2-7B-Base', 'baichuan-inc/Baichuan2-7B-Base'),
                Model('baichuan-inc/Baichuan2-13B-Chat', 'baichuan-inc/Baichuan2-13B-Chat'),
                Model('baichuan-inc/Baichuan2-13B-Base', 'baichuan-inc/Baichuan2-13B-Base'),
            ]),
            ModelGroup([
                Model('baichuan-inc/Baichuan2-7B-Chat-4bits', 'baichuan-inc/Baichuan2-7B-Chat-4bits'),
                Model('baichuan-inc/Baichuan2-13B-Chat-4bits', 'baichuan-inc/Baichuan2-13B-Chat-4bits'),
            ],
                       requires=['bitsandbytes<0.41.2', 'accelerate<0.26'])
        ],
        Baichuan2Loader,
        template=TemplateType.baichuan,
        architectures=['BaichuanForCausalLM', 'BaiChuanForCausalLM'],
        model_arch=ModelArch.baichuan,
    ))