# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import gc
import os
import re
import warnings
from contextlib import contextmanager
from functools import partial
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Type, Union

import numpy as np
import paddle
import paddle.nn as nn
from paddle import Tensor
from paddle.distributed.fleet.meta_parallel.parallel_layers import PipelineLayer

try:
    from paddle.distributed.fleet.meta_parallel import LocalSharedLayerDesc
except:
    LocalSharedLayerDesc = None
from paddle.nn import Layer

from ......utils import logging
from ......utils.deps import is_dep_available, require_deps
from ...tokenizer.tokenizer_utils import InitTrackerMeta, adapt_stale_fwd_patch
from ..generation import GenerationConfig, GenerationMixin
from ..utils import (
    ASYMMETRY_QUANT_SCALE_MAX,
    ASYMMETRY_QUANT_SCALE_MIN,
    CONFIG_NAME,
    LEGACY_CONFIG_NAME,
    PADDLE_WEIGHTS_INDEX_NAME,
    PADDLE_WEIGHTS_NAME,
    PYTORCH_WEIGHTS_INDEX_NAME,
    PYTORCH_WEIGHTS_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SYMMETRY_QUANT_SCALE,
    device_guard,
    resolve_file_path,
)
from .configuration_utils import PretrainedConfig
from .conversion_utils import ConversionMixin
from .utils import (
    ContextManagers,
    fn_args_to_dict,
    get_checkpoint_shard_files,
    paddlenlp_load,
    weight_name_suffix,
)

__all__ = [
    "PretrainedModel",
]


def _add_variant(weights_name: str, variant=None) -> str:
    if variant is not None and len(variant) > 0:
        splits = weights_name.split(".")
        splits = splits[:-1] + [variant] + splits[-1:]
        weights_name = ".".join(splits)

    return weights_name


@contextmanager
def dtype_guard(dtype="float32"):
    origin_dtype = paddle.get_default_dtype()
    paddle.set_default_dtype(dtype)
    try:
        yield
    finally:
        paddle.set_default_dtype(origin_dtype)


_init_weights = True


@contextmanager
def no_init_weights(_enable=True):
    """
    Context manager to globally disable weight initialization to speed up loading large models.

    TODO(Patrick): Delete safety argument `_enable=True` at next major version. .
    """
    global _init_weights
    old_init_weights = _init_weights
    if _enable:
        _init_weights = False
    try:
        yield
    finally:
        _init_weights = old_init_weights


def _split_keys_evenly(keys: list, n: int) -> list:
    """Split a list into n lists with an equal number of elements.

    Args:
        keys (list): the list to be split
        n (int): number of splits

    Returns:
        result: list of lists
    """

    total_len = len(keys)
    base_size = total_len // n
    extra = total_len % n

    result = []
    index = 0
    for _ in range(n):
        part_size = base_size + 1 if extra > 0 else base_size
        extra -= 1
        result.append(keys[index : index + part_size])
        index += part_size

    return result


def _load_part_state_dict_from_safetensors(
    keys,
    checkpoint_file: Union[str, os.PathLike],
    tensor_parallel_split_mapping,
    fliter_dict_keys,
    device,
    quantization_linear_list=None,
    quantization_config=None,
    dtype=None,
    return_numpy=False,
    convert_from_hf=False,
    transpose_weight_keys=None,
):
    import paddle
    from safetensors import safe_open

    if transpose_weight_keys:
        transpose_weight_keys = set(transpose_weight_keys)

    def _is_need_transpose(key):
        if "lora" not in key and convert_from_hf and transpose_weight_keys:
            return key in transpose_weight_keys

    def _transpose_hf_weight(key, weight):
        if _is_need_transpose(key):
            return weight.transpose([-1, -2])
        return weight

    part_state_dict = {}
    scale_dict = {}
    with safe_open(checkpoint_file, framework="paddle") as f:
        for key in keys:
            # 1. non-merge ckpt loading dont have filter key.
            # 2. merge ckpt will skip quant scale by `fliter_dict_keys`
            if (
                key.endswith(SYMMETRY_QUANT_SCALE)
                or key.endswith(ASYMMETRY_QUANT_SCALE_MIN)
                or key.endswith(ASYMMETRY_QUANT_SCALE_MAX)
            ):
                continue

            if fliter_dict_keys is not None and key not in fliter_dict_keys:
                continue

            py_safe_slice_ = f.get_slice(key)
            if (
                quantization_linear_list is not None
                and key.split(".weight")[0] in quantization_linear_list
                and not key.endswith("_scale")
            ):
                raise NotImplementedError
            else:
                if key in tensor_parallel_split_mapping:
                    tp_fn = tensor_parallel_split_mapping[key]
                    if _is_need_transpose(key):
                        assert isinstance(tp_fn, partial)
                        is_column = True
                        if "is_column" in tp_fn.keywords:
                            is_column = tp_fn.keywords["is_column"]
                        is_column = not is_column
                        tp_fn = partial(
                            tp_fn.func,
                            *tp_fn.args,
                            **{**tp_fn.keywords, "is_column": is_column},
                        )
                    if len(py_safe_slice_.shape) == 0:
                        weight = tp_fn(py_safe_slice_[:])
                    else:
                        weight = tp_fn(py_safe_slice_)
                else:
                    weight = py_safe_slice_[:]

                if not return_numpy and device == "expected":
                    weight = weight._copy_to(
                        paddle.framework._current_expected_place(),
                        True,
                    )
                weight = _transpose_hf_weight(key, weight)
                if return_numpy:
                    weight = weight.numpy()
                part_state_dict[key] = weight

        for key in keys:
            if (
                key.endswith(SYMMETRY_QUANT_SCALE)
                or key.endswith(ASYMMETRY_QUANT_SCALE_MIN)
                or key.endswith(ASYMMETRY_QUANT_SCALE_MAX)
            ):
                scale = f.get_tensor(key)
                if not return_numpy and device == "expected":
                    scale = scale._copy_to(
                        paddle.framework._current_expected_place(), False
                    )
                if return_numpy:
                    scale = scale.numpy()
                scale_dict[key] = scale
    return part_state_dict, scale_dict


def load_state_dict(
    checkpoint_file: Union[str, os.PathLike],
    tensor_parallel_split_mapping=None,
    fliter_dict_keys=None,
    device="cpu",
    ckpt_quant_stage="O0",
    convert_from_hf=False,
    transpose_weight_keys=None,
):
    """
    Reads a PaddlePaddle checkpoint file, returning properly formatted errors if they arise.
    """

    if tensor_parallel_split_mapping is None:
        tensor_parallel_split_mapping = {}

    if Path(checkpoint_file).suffix == ".safetensors":
        require_deps("safetensors")
        from safetensors import safe_open

        with safe_open(checkpoint_file, framework="paddle") as f:
            keys = list(f.keys())
        state_dict, scale_dict = _load_part_state_dict_from_safetensors(
            keys,
            checkpoint_file,
            tensor_parallel_split_mapping,
            fliter_dict_keys,
            "expected",
            dtype=None,
            return_numpy=False,
            convert_from_hf=convert_from_hf,
            transpose_weight_keys=transpose_weight_keys,
        )
    else:
        state_dict = paddlenlp_load(checkpoint_file, map_location="cpu")
    return state_dict


_re_layer_prefix = re.compile(r"\.(\d+)\.")


def _load_state_dict_into_model(
    model_to_load, state_dict, start_prefix, convert_from_hf
):
    # torch will cast dtype in load_state_dict, but paddle strictly check dtype
    _convert_state_dict_dtype_and_shape(state_dict, model_to_load, convert_from_hf)

    error_msgs = []

    if len(start_prefix) > 0:
        for key in list(state_dict.keys()):
            if key.startswith(start_prefix):
                state_dict[key.replace(start_prefix, "")] = state_dict.pop(key)

    # TODO: add return status to state_dict
    with warnings.catch_warnings(record=True) as w:
        warnings.resetwarnings()
        # paddlenlp hold  missing_keys , just ignore not found warnings.
        warnings.filterwarnings(
            "ignore", message=r".*is not found in the provided dict.*"
        )
        warnings.filterwarnings("ignore", message=r".*paddle.to_tensor.*")
        if convert_from_hf:
            try:
                model_to_load.set_hf_state_dict(state_dict)
            except NotImplementedError:
                pass
        model_to_load.set_state_dict(state_dict)
        error_msgs.extend([str(x.message) for x in w])

    del state_dict

    return error_msgs


def _convert_state_dict_dtype_and_shape(state_dict, model_to_load, convert_from_hf):
    # convert the dtype of state dict
    def is_0d_or_1d(tensor):
        return len(tensor.shape) == 0 or list(tensor.shape) == [1]

    if convert_from_hf:
        model_state_dict = model_to_load.get_hf_state_dict()
    else:
        model_state_dict = model_to_load.state_dict()
    for key, value in model_state_dict.items():
        if key in list(state_dict.keys()):
            if isinstance(state_dict[key], np.ndarray):
                raise ValueError(
                    "convert_state_dict_dtype expected paddle.Tensor not numpy.ndarray, please convert numpy.ndarray to paddle.Tensor"
                )
            # confirm parameter cast is executed on the same device as model
            # TODO: cast(FP32 -> FP16) has diff on different devices, need to fix it
            if (
                state_dict[key].is_floating_point()
                and state_dict[key].dtype != value.dtype
            ):
                state_dict[key] = paddle.cast(state_dict.pop(key), value.dtype)
            # unified 0d and 1d tensor
            if is_0d_or_1d(value) and is_0d_or_1d(state_dict[key]):
                if list(value.shape) != list(state_dict[key].shape):
                    state_dict[key] = paddle.reshape(state_dict.pop(key), value.shape)


def _load_state_dict_into_meta_model(
    model,
    state_dict,
    loaded_state_dict_keys,  # left for now but could be removed, see below
    start_prefix,
    expected_keys,
    dtype=None,
    is_safetensors=False,
    keep_in_fp32_modules=None,
):
    """
    This is somewhat similar to `_load_state_dict_into_model`, but deals with a model that has some or all of its
    params on a `meta` device. It replaces the model params with the data from the `state_dict`, while moving the
    params back to the normal device, but only for `loaded_state_dict_keys`.

    `start_prefix` is used for models which insert their name into model keys, e.g. `bert` in
    `bert.pooler.dense.weight`

    """
    from paddle.common_ops_import import convert_np_dtype_to_dtype_

    dtype = convert_np_dtype_to_dtype_(dtype)
    error_msgs = []
    model_state_dict = model.state_dict()
    for param_name, param in state_dict.items():
        # First part of the test is always true as loaded_state_dict_keys always contains state_dict keys.
        if param_name not in loaded_state_dict_keys or param_name not in expected_keys:
            continue

        if param_name.startswith(start_prefix):
            param_name = param_name[len(start_prefix) :]

        if param.place != paddle.framework._current_expected_place():
            param = param._copy_to(paddle.framework._current_expected_place(), False)

        # # We convert floating dtypes to the `dtype` passed. We want to keep the buffers/params
        # # in int/uint/bool and not cast them.
        if dtype is not None and paddle.is_floating_point(param):
            if (
                keep_in_fp32_modules is not None
                and any(
                    module_to_keep_in_fp32 in param_name
                    for module_to_keep_in_fp32 in keep_in_fp32_modules
                )
                and (dtype == paddle.float16 or dtype == paddle.bfloat16)
            ):
                param = param.astype(dtype=paddle.float32)
            else:
                param = param.astype(dtype=dtype)

        if dtype is None:
            old_param = model
            splits = param_name.split(".")
            for split in splits:
                old_param = getattr(old_param, split)
                if old_param is None:
                    break

            if old_param is not None:
                param = param.astype(dtype=old_param.dtype)
        with paddle.no_grad():
            model_state_dict[param_name].get_tensor()._share_data_with(
                param.value().get_tensor()
            )
            param.value().get_tensor()._clear()
    return error_msgs


class PretrainedModel(
    Layer, GenerationMixin, ConversionMixin, metaclass=InitTrackerMeta
):
    """
    The base class for all pretrained models. It mainly provides common methods
    for loading (construction and loading) and saving pretrained models. Loading
    and saving also rely on the following class attributes which should be overridden
    by derived classes accordingly:

    - **model_config_file** (str): Represents the file name of model configuration
        for configuration saving and loading in local file system. The value is
        `model_config.json`.
    - **resource_files_names** (dict): Name of local file where the model configuration
        can be saved and loaded locally. Currently, resources only include the model state,
        thus the dict only includes `'model_state'` as key with corresponding
        value `'model_state.pdparams'` for model weights saving and loading.
    - **pretrained_init_configuration** (dict): Provides the model configurations
        of built-in pretrained models (contrasts to models in local file system).
        It has pretrained model names as keys (such as `bert-base-uncased`), and
        the values are dict preserving corresponding configuration for model initialization.
    - **pretrained_resource_files_map** (dict): Provides resource URLs of built-in
        pretrained models (contrasts to models in local file system).
        It has the same key as resource_files_names (that is "model_state"),
        and the corresponding value is a dict with specific model name to model weights URL mapping
        (such as "bert-base-uncased" ->
        "https://bj.bcebos.com/paddlenlp/models/transformers/bert-base-uncased.pdparams").
    - **base_model_prefix** (str): Represents the attribute associated to the
        base model in derived classes of the same architecture adding layers on
        top of the base model. Note: A base model class is pretrained model class
        decorated by `register_base_model`, such as `BertModel`; A derived model
        class is a pretrained model class adding layers on top of the base model,
        and it has a base model as attribute, such as `BertForSequenceClassification`.

    Methods common to models for text generation are defined in `GenerationMixin`
    and also inherited here.

    Besides, metaclass `InitTrackerMeta` is used to create `PretrainedModel`,
    by which subclasses can track arguments for initialization automatically.
    """

    # Deprecated(wj-Mcat): after 2.6.* version
    # save the old-school `LEGACY_CONFIG_NAME`, and will be changed to `CONFIG_NAME` after 2.6.* version
    model_config_file = LEGACY_CONFIG_NAME

    pretrained_init_configuration = {}
    # TODO: more flexible resource handle, namedtuple with fields as:
    # resource_name, saved_file, handle_name_for_load(None for used as __init__
    # arguments), handle_name_for_save
    resource_files_names = {"model_state": PADDLE_WEIGHTS_NAME}
    pretrained_resource_files_map = {}
    base_model_prefix = ""
    main_input_name = "input_ids"
    config_class = None
    _keep_in_fp32_modules = None

    # a list of `re` patterns of `state_dict` keys that should be removed from the list of missing
    # keys we find (keys inside the model but not in the checkpoint) and avoid unnecessary warnings.
    _keys_to_ignore_on_load_missing = None
    # a list of `re` patterns of `state_dict` keys that should be removed from the list of
    # unexpected keys we find (keys inside the checkpoint but not the model) and avoid unnecessary
    # warnings.
    _keys_to_ignore_on_load_unexpected = None
    # a list of `state_dict` keys to ignore when saving the model (useful for keys that aren't
    # trained, but which are either deterministic or tied variables)
    _keys_to_ignore_on_save = None
    _tied_weights_keys = None

    def __init__(self, *args, **kwargs):
        super(PretrainedModel, self).__init__()

        if not self.constructed_from_pretrained_config():
            return

        # extract config from args
        config = None
        for arg in args:
            if isinstance(arg, PretrainedConfig):
                config = arg
                break
        if config is not None:
            self.config: PretrainedConfig = config
            self.model_config_file = CONFIG_NAME
            self.generation_config = (
                GenerationConfig.from_model_config(self.config)
                if self.can_generate()
                else None
            )
            return

        # extract config from kwargs
        if "config" not in kwargs:
            raise ValueError(
                "PretrainedConfig instance not found in the arguments, you can set it as args or kwargs with config field"
            )

        config = kwargs["config"]
        if not isinstance(config, PretrainedConfig):
            raise TypeError(
                "config parameter should be the instance of PretrainedConfig"
            )

        self.config: PretrainedConfig = kwargs["config"]
        self.generation_config = (
            GenerationConfig.from_model_config(self.config)
            if self.can_generate()
            else None
        )
        self.model_config_file = CONFIG_NAME
        self.warnings_issued = {}

    def _post_init(self, original_init, *args, **kwargs):
        """
        It would be hooked after `__init__` to add a dict including arguments of
        `__init__` as a attribute named `config` of the pretrained model instance.
        """
        if not self.constructed_from_pretrained_config():
            init_dict = fn_args_to_dict(original_init, *((self,) + args), **kwargs)
            self.config = init_dict

        # only execute when it's the base method
        if (
            original_init.__module__ != "paddlenlp.transformers.model_utils"
            and self.__class__.init_weights is PretrainedModel.init_weights
        ):
            self.init_weights()

        # Note:
        # 1. PipelineLayer will create parameters for each layer and
        # call `_synchronize_shared_weights()` to synchronize the shared parameters.
        # 2. When setting the model `state_dict`, `_synchronize_shared_weights` will be called to
        # synchronize the shared parameters.
        # However, `self._init_weights` will re-initialize the parameters without
        # synchronizing the shared parameters. If the following step does not load a checkpoint,
        # the shared parameters will be different.

        if isinstance(self, PipelineLayer):
            self._synchronize_shared_weights()

    def _init_weights(self, layer):
        """
        Initialize the weights. This method should be overridden by derived class.
        """
        pass

    def _initialize_weights(self, layer):
        """
        Initialize the weights if they are not already initialized.
        """
        if getattr(layer, "_is_initialized", False):
            return
        self._init_weights(layer)
        layer._is_initialized = True

    def init_weights(self):
        """
        If needed prunes and maybe initializes weights. If using a custom `PreTrainedModel`, you need to implement any
        initialization logic in `_init_weights`.
        """
        # call pure
        if _init_weights:
            # Initialize weights
            self.apply(self._initialize_weights)

            # Tie weights should be skipped when not initializing all weights
            # since from_pretrained(...) calls tie weights anyways

            # TODO(wj-Mcat): enable all tie-weights later
            # self.tie_weights()

    @classmethod
    def _from_config(cls, config, **kwargs):
        """
        All context managers that the model should be initialized under go here.

        Args:
            dtype (`paddle.dtype`, *optional*):
                Override the default `paddle.dtype` and load the model under this dtype.
        """
        dtype = kwargs.pop("dtype", None)

        if dtype is None:
            if config.dtype is not None:
                dtype = config.dtype
            else:
                dtype = paddle.get_default_dtype()

        with dtype_guard(dtype):
            model = cls(config, **kwargs)

        return model

    @classmethod
    def from_config(cls, config, **kwargs):
        """
        All context managers that the model should be initialized under go here.

        Args:
            dtype (`paddle.dtype`, *optional*):
                Override the default `paddle.dtype` and load the model under this dtype.
        """
        return cls._from_config(config, **kwargs)

    @classmethod
    def set_inference_config(cls, config, predictor_args, **kwargs):
        """
        All inference config can set here.
        Args:
            config : PretrainedConfig
                The config of the model.
            predictor_args : PredictorArgument
                The args of the predictor.
        """
        tensor_parallel_degree = kwargs.pop("tensor_parallel_degree", 1)
        tensor_parallel_rank = kwargs.pop("tensor_parallel_rank", 0)

        if predictor_args.mode == "dynamic" or predictor_args.speculate_method in [
            "eagle",
            "mtp",
        ]:
            config.tensor_parallel_degree = tensor_parallel_degree
            config.tensor_parallel_rank = tensor_parallel_rank
            config.model_name_or_path = predictor_args.model_name_or_path
            config.quant_type = predictor_args.quant_type
            config.cachekv_int8_type = predictor_args.cachekv_int8_type
            config.use_fake_parameter = predictor_args.use_fake_parameter
            config.single_card_ptq = not predictor_args.use_fake_parameter
        config.append_attn = predictor_args.append_attn
        config.decode_strategy = predictor_args.decode_strategy
        config.mla_use_matrix_absorption = predictor_args.mla_use_matrix_absorption
        config.weightonly_group_size = predictor_args.weightonly_group_size
        config.weight_block_size = predictor_args.weight_block_size
        config.moe_quant_type = predictor_args.moe_quant_type

        if predictor_args.block_attn:
            config.block_size = predictor_args.block_size
            config.max_seq_len = predictor_args.total_max_length

        if predictor_args.speculate_method is not None:
            config.speculate_method = predictor_args.speculate_method
            config.speculate_max_draft_token_num = (
                predictor_args.speculate_max_draft_token_num
            )
            config.speculate_verify_window = predictor_args.speculate_verify_window
            config.speculate_max_candidate_len = (
                predictor_args.speculate_max_candidate_len
            )
            if predictor_args.speculate_method == "inference_with_reference":
                config.speculate_max_ngram_size = (
                    predictor_args.speculate_max_ngram_size
                )
            if predictor_args.speculate_method is not None:
                if not config.get("speculate_model_type", "None") in ["eagle", "mtp"]:
                    config.decode_strategy = "speculate_decoding"
        config.return_full_hidden_states = predictor_args.return_full_hidden_states

    @classmethod
    def confirm_inference_model(cls, predictor_args, **kwargs):
        """
        Confirm the inference model whether it need to change the AVX inference Model
        Args:
            model : PretrainedModel
                The model for inference.
            predictor_args : PredictorArgument
                The args of the predictor.
        """
        return cls

    @property
    def base_model(self):
        """
        PretrainedModel: The body of the same model architecture. It is the base
            model itself for base model or the base model attribute for derived
            model.
        """
        return getattr(self, self.base_model_prefix, self)

    @property
    def model_name_list(self):
        """
        list: Contains all supported built-in pretrained model names of the
            current PretrainedModel class.
        """
        # Todo: return all model name
        return list(self.pretrained_init_configuration.keys())

    def can_generate(self) -> bool:
        """
        Returns whether this model can generate sequences with `.generate()`.
        Returns:
            `bool`: Whether this model can generate sequences with `.generate()`.
        """
        # Detects whether `prepare_inputs_for_generation` has been overwritten, which is a requirement for generation
        if "GenerationMixin" in str(self.prepare_inputs_for_generation):
            return False
        return True

    def recompute_enable(self):
        r"""
        Enable Recompute.
        All layers with the `enable_recompute` attribute will be set to `True`
        """

        def fn(layer):
            if hasattr(layer, "enable_recompute") and (
                layer.enable_recompute is False or layer.enable_recompute == 0
            ):
                layer.enable_recompute = True

        self.apply(fn)

    def recompute_disable(self):
        r"""
        Disable Recompute.
        All layers with the `enable_recompute` attribute will be set to `False`
        """

        def fn(layer):
            if hasattr(layer, "enable_recompute") and (
                layer.enable_recompute is False or layer.enable_recompute == 0
            ):
                layer.enable_recompute = True

        self.apply(fn)

    def tie_weights(self):
        """
        Tie the weights between the input embeddings and the output embeddings.
        """
        if self.config.tie_word_embeddings:
            output_embeddings = self.get_output_embeddings()
            input_embeddings = self.get_input_embeddings()
            if output_embeddings is not None and input_embeddings is not None:
                if input_embeddings.weight.shape != output_embeddings.weight.shape:
                    logging.warning(
                        f"The shape of input embeddings is {input_embeddings.weight.shape} and the shape of output embeddings is {output_embeddings.weight.shape}. "
                        "This is only expected if you are calling the `resize_token_embeddings` method"
                    )
                output_embeddings.weight = input_embeddings.weight
                if getattr(output_embeddings, "bias", None) is not None:
                    # need to pad
                    if (
                        output_embeddings.weight.shape[0]
                        > output_embeddings.bias.shape[0]
                    ):
                        old_bias = output_embeddings.bias
                        pad_length = (
                            output_embeddings.weight.shape[0] - old_bias.shape[0]
                        )
                        output_embeddings.bias = output_embeddings.create_parameter(
                            shape=[output_embeddings.weight.shape[0]],
                            attr=output_embeddings._bias_attr,
                            dtype=output_embeddings._dtype,
                            is_bias=True,
                        )
                        new_bias = paddle.concat(
                            [
                                old_bias,
                                paddle.zeros(
                                    [pad_length], dtype=output_embeddings.bias.dtype
                                ),
                            ]
                        )
                        output_embeddings.bias.set_value(new_bias)
                    # need to trim
                    elif (
                        output_embeddings.weight.shape[0]
                        < output_embeddings.bias.shape[0]
                    ):
                        new_bias = output_embeddings.bias[
                            : output_embeddings.weight.shape[0]
                        ]
                        output_embeddings.bias = output_embeddings.create_parameter(
                            shape=[output_embeddings.weight.shape[0]],
                            attr=output_embeddings._bias_attr,
                            dtype=output_embeddings._dtype,
                            is_bias=True,
                        )
                        output_embeddings.bias.set_value(new_bias)

    def resize_position_embeddings(self, new_num_position_embeddings: int):
        """resize position embedding, this method should be overrited overwrited by downstream models

        Args:
            new_num_position_embeddings (int): the new position size

        Raises:
            NotImplementedError: when called and not be implemented
        """
        raise NotImplementedError(
            f"`resize_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__} in `{self.__class__.__module__}.py`"
        )

    @classmethod
    def constructed_from_pretrained_config(cls, init_func=None) -> bool:
        """check if the model is constructed from `PretrainedConfig`
        Returns:
            bool: if the model is constructed from `PretrainedConfig`
        """
        return cls.config_class is not None and issubclass(
            cls.config_class, PretrainedConfig
        )

    def resize_token_embeddings(
        self, new_num_tokens: Optional[int] = None
    ) -> nn.Embedding:
        """
        Resizes input token embeddings matrix of the model according to new_num_tokens.

        Args:
            new_num_tokens (Optional[int]):
                The number of new tokens in the embedding matrix. Increasing the size will add newly initialized
                vectors at the end. Reducing the size will remove vectors from the end. If not provided or None, just
                returns a pointer to the input tokens embedding module of the model without doing anything.

        Returns:
            paddle.nn.Embedding: The input tokens Embeddings Module of the model.
        """
        old_embeddings: nn.Embedding = self.get_input_embeddings()
        if not new_num_tokens or new_num_tokens == old_embeddings.weight.shape[0]:
            return old_embeddings

        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
        self.set_input_embeddings(new_embeddings)

        # 2. Update vocab_size
        self.base_model.config["vocab_size"] = new_num_tokens
        self.vocab_size = new_num_tokens

        # update init_config
        self._update_init_config(self.init_config, "vocab_size", new_num_tokens)

        # Tie the weights between the input embeddings and the output embeddings if needed.
        self.tie_weights()

        return new_embeddings

    def _update_init_config(self, init_config: dict, key: str, value: Any):
        """update init_config by <key, value> pair

        Args:
            init_config (dict): the init_config instance
            key (str): the key field
            value (Any): the new value of instance
        """
