modeling_utils

mindnlp.transformers.modeling_utils.PreTrainedModel

Bases: Module, CellUtilMixin, GenerationMixin, PeftAdapterMixin

Abstract class for Pretrained models

Source code in mindnlp/transformers/modeling_utils.py

class PreTrainedModel(nn.Module, CellUtilMixin, GenerationMixin, PeftAdapterMixin):
    """
    Abstract class for Pretrained models
    """
    config_class = None
    base_model_prefix = ""
    main_input_name = "input_ids"

    # 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
    _keys_to_ignore_on_save = None

    _tied_weights_keys = None

    _keep_in_fp32_modules = None

    supports_recompute = False

    def __init__(self, config):
        """

        Initializes a new instance of the PreTrainedModel class.

        Args:
            self: The instance of the class.
            config: A dictionary containing the configuration parameters for the model.

        Returns:
            None.

        Raises:
            ValueError: If the config parameter is invalid or missing required fields.
            TypeError: If the config parameter is not of the expected type.
            RuntimeError: If an error occurs during the initialization process.
        """
        super().__init__()
        self._check_and_unset_acl()
        # Save config in model
        self.config = config
        self.name_or_path = config.name_or_path
        self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None

    def _check_and_unset_acl(self):
        """
        This method '_check_and_unset_acl' is defined within the class 'PreTrainedModel'
        and is used to verify and remove the ACL (Access Control List) settings if certain conditions are met.

        Args:
            self: An instance of the class 'PreTrainedModel'.
                This parameter is used to access the attributes and methods of the class instance.

        Returns:
            None.

        Raises:
            None.
        """
        if "MS" in str(self.__class__.__name__) and \
            'MS_DEV_FORCE_ACL' in os.environ:
            del os.environ['MS_DEV_FORCE_ACL']

    def post_init(self):
        """
        A method executed at the end of each Transformer model initialization, to execute code that needs the model's
        modules properly initialized (such as weight initialization).

        """
        self.init_weights()

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

        Args:
            torch_dtype (`torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model under this dtype.
        """
        model = cls(config, **kwargs)

        return model

    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`.
        """
        # Prune heads if needed
        if self.config.pruned_heads:
            self.prune_heads(self.config.pruned_heads)

        if _init_weights:
            # Initialize weights
            if getattr(self, 'apply', None):
                self.apply(self._initialize_weights)
            else:
                for _, cell in self.name_cells().items():
                    self._initialize_weights(cell)

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

    def prune_heads(self, heads_to_prune: Dict[int, List[int]]):
        """
        Prunes heads of the base model.

        Arguments:
            heads_to_prune (`Dict[int, List[int]]`):
                Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads
                to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on
                layer 1 and heads 2 and 3 on layer 2.
        """
        # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
        for layer, heads in heads_to_prune.items():
            union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)
            self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSON

        self.base_model._prune_heads(heads_to_prune)

    def _init_weights(self, cell):
        """
        Initialize the weights. This method should be overridden by derived class and is
        the only initialization method that will be called when loading a checkpoint
        using `from_pretrained`. Any attempt to initialize outside of this function
        will be useless as the torch.nn.init function are all replaced with skip.
        """
    def _initialize_weights(self, module):
        """
        Initialize the weights if they are not already initialized.
        """
        if getattr(module, "_is_initialized", False):
            return
        self._init_weights(module)
        module._is_initialized = True

    @property
    def base_model(self):
        """
        to get base_model
        """
        return getattr(self, self.base_model_prefix, self)

    def get_input_embeddings(self) -> "nn.Module":
        """
        Returns the model's input embeddings.

        Returns:
            :obj:`nn.Module`: A mindspore cell mapping vocabulary to hidden states.
        """
        base_model = getattr(self, self.base_model_prefix, self)
        print(base_model)
        if base_model is not self:
            return base_model.get_input_embeddings()
        raise NotImplementedError

    def set_input_embeddings(self, new_embeddings: nn.Module):
        """
        Set model's input embeddings.

        Args:
            value (:obj:`nn.Module`): A mindspore cell mapping vocabulary to hidden states.
        """
        base_model = getattr(self, self.base_model_prefix, self)
        if base_model is not self:
            return base_model.set_input_embeddings(new_embeddings)
        raise NotImplementedError

    def resize_position_embeddings(self, new_num_position_embeddings: int):
        """
        resize the model position embeddings if necessary
        """
        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__}"
        )

    def get_output_embeddings(self):
        """ Get model's output embeddings
            Return None if the model doesn't have output embeddings
        """
        return None  # Overwrite for models with output embeddings

    def set_output_embeddings(self, new_embeddings: nn.Module):
        """
        Set model's output embeddings.

        Args:
            value (:obj:`nn.Module`): A mindspore cell mapping vocabulary to hidden states.
        """
        base_model = getattr(self, self.base_model_prefix, self)
        if base_model is not self:
            return base_model.set_output_embeddings(new_embeddings)
        raise NotImplementedError

    def get_position_embeddings(self):
        """
        get the model position embeddings if necessary
        """
        raise NotImplementedError(
            f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__}"
        )

    def tie_weights(self):
        """
        Make sure we are sharing the input and output embeddings.
        If you need this feature,
        you need to get it yourself output Add the output you need to add to the embeddings function_ Embedding layer,
        otherwise you cannot
        """
        if getattr(self.config, "tie_word_embeddings", True):
            output_embeddings = self.get_output_embeddings() # pylint: disable=assignment-from-none
            if output_embeddings is not None:
                self._tie_or_clone_weights(
                    output_embeddings, self.get_input_embeddings())

        if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False):
            if hasattr(self, self.base_model_prefix):
                self = getattr(self, self.base_model_prefix) # pylint: disable=self-cls-assignment
            self._tie_encoder_decoder_weights(
                self.encoder, self.decoder, self.base_model_prefix)

        for _, cell in self.cells_and_names():
            if hasattr(cell, "_tie_weights"):
                cell._tie_weights()

    @staticmethod
    def _tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str):
        """tie encoder decoder weights"""
        uninitialized_encoder_weights: List[str] = []
        if decoder.__class__ != encoder.__class__:
            logger.info(
                f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder"
                " weights are correctly initialized."
            )

        def tie_encoder_to_decoder_recursively(
            decoder_pointer: nn.Module,
            encoder_pointer: nn.Module,
            module_name: str,
            uninitialized_encoder_weights: List[str],
            depth=0,
        ):
            assert isinstance(decoder_pointer, nn.Module) and isinstance(
                encoder_pointer, nn.Module
            ), f"{decoder_pointer} and {encoder_pointer} have to be of type nn.Module"
            if hasattr(decoder_pointer, "weight"):
                assert hasattr(encoder_pointer, "weight")
                encoder_pointer.weight = decoder_pointer.weight
                encoder_pointer._params['weight'] = decoder_pointer.weight
                if hasattr(decoder_pointer, "bias"):
                    assert hasattr(encoder_pointer, "bias")
                    encoder_pointer.bias = decoder_pointer.bias
                    encoder_pointer._params['bias'] = decoder_pointer.bias
                return

            encoder_cells = encoder_pointer._cells
            decoder_cells = decoder_pointer._cells
            if len(decoder_cells) > 0:
                assert (
                    len(encoder_cells) > 0
                ), f"Encoder cell {encoder_pointer} does not match decoder cell {decoder_pointer}"

                all_encoder_weights = {module_name + "/" + sub_name for sub_name in encoder_cells.keys()}
                encoder_layer_pos = 0
                for name, _ in decoder_cells.items():
                    if name.isdigit():
                        encoder_name = str(int(name) + encoder_layer_pos)
                        decoder_name = name
                        if not isinstance(decoder_cells[decoder_name], type(encoder_cells[encoder_name])) and len(
                            encoder_cells
                        ) != len(decoder_cells):
                            # this can happen if the name corresponds to the position in a list module list of layers
                            # in this case the decoder has added a cross-attention that the encoder does not have
                            # thus skip this step and subtract one layer pos from encoder
                            encoder_layer_pos -= 1
                            continue
                    elif name not in encoder_cells:
                        continue
                    elif depth > 500:
                        raise ValueError(
                            "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is"
                            " a circular dependency between two or more `nn.Module` of your model."
                        )
                    else:
                        decoder_name = encoder_name = name
                    tie_encoder_to_decoder_recursively(
                        decoder_cells[decoder_name],
                        encoder_cells[encoder_name],
                        module_name + "/" + name,
                        uninitialized_encoder_weights,
                        depth=depth + 1,
                    )
                    all_encoder_weights.remove(module_name + "/" + encoder_name)

                uninitialized_encoder_weights += list(all_encoder_weights)

        # tie weights recursively
        tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights)
        if len(uninitialized_encoder_weights) > 0:
            logger.warning(
                f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}"
            )

    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
        """ Tie or clone module weights depending of weither we are using or not
        """
        if hasattr(output_embeddings, 'weight'):
            output_embeddings.weight = input_embeddings.weight

        if getattr(output_embeddings, "bias", None) is not None:
            if output_embeddings.weight.shape[0] == output_embeddings.bias.shape[0]:
                pass
            else:
                # instantial a new Parameter since mindspore.Parameter do not support assign_value with different shape
                if output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0] > 0:
                    new_bias = F.pad(
                        output_embeddings.bias.data,
                        (0, output_embeddings.weight.shape[0] -
                        output_embeddings.bias.shape[0]),
                        "constant",
                        0,
                    )
                else:
                    new_bias = output_embeddings.bias[:output_embeddings.weight.shape[0]]
                new_bias = Parameter(new_bias, name=output_embeddings.bias.name, requires_grad=output_embeddings.bias.requires_grad)
                output_embeddings.bias = new_bias

        if hasattr(output_embeddings, "out_channels") and hasattr(input_embeddings, "vocab_size"):
            output_embeddings.out_channels = input_embeddings.vocab_size

    def resize_token_embeddings(
        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
    ) -> nn.Embedding:
        """
        Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.

        Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.

        Arguments:
            new_num_tokens (`int`, *optional*):
                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 `torch.nn.Embedding` module of the model without doing anything.

        Returns:
            `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
        """
        model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
        if new_num_tokens is None and pad_to_multiple_of is None:
            return model_embeds
        # Update base model and current model config
        self.config.vocab_size = model_embeds.weight.shape[0]
        self.vocab_size = model_embeds.weight.shape[0]
        # Tie weights again if needed
        self.tie_weights()

        return model_embeds

    def _resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None):
        """
        Resize the token embeddings of the PreTrainedModel.

        Args:
            self (PreTrainedModel): The instance of the PreTrainedModel class.
            new_num_tokens (int): The desired number of tokens for the resized embeddings.
            pad_to_multiple_of (int, optional): The value to which the number of tokens should be padded. 
                Defaults to None.

        Returns:
            None: The method modifies the input and output embeddings of the model in-place.

        Raises:
            None.
        """
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens, pad_to_multiple_of)

        self.set_input_embeddings(new_embeddings)
        # self.get_input_embeddings().weight.data_sync(True)
        # Update new_num_tokens with the actual size of new_embeddings
        if pad_to_multiple_of is not None:
            new_num_tokens = new_embeddings.weight.shape[0]
        # if word embeddings are not tied, make sure that lm head is resized as well
        if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings:
            old_lm_head = self.get_output_embeddings() # pylint: disable=assignment-from-none
            new_lm_head = self._get_resized_lm_head(
                old_lm_head, new_num_tokens)
            self.set_output_embeddings(new_lm_head)
            self.get_output_embeddings().weight.data_sync(True)

        return self.get_input_embeddings()

    def resize_tokenizer_embeddings(self, new_num_tokens):
        """
        Obtain a new embedding layer or use the original one without updating it.
        """
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(
            old_embeddings, new_num_tokens)
        self.set_input_embeddings(new_embeddings)
        return self.get_input_embeddings()

    def _get_resized_embeddings(
        self,
        old_embeddings: nn.Embedding,
        new_num_tokens: Optional[int] = None,
        pad_to_multiple_of: Optional[int] = None,
    ) -> nn.Embedding:
        """
        Build a resized Embedding Module from a provided token Embedding Module.
        Increasing the size will add newly initialized vectors at the end
        Reducing the size will remove vectors from the end

        Args:
            new_num_tokens: (`optional`) int
                New number of 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: return the provided token Embedding Module.

        Returns: ``mindspore.nn.Embeddings``
            Pointer to the resized Embedding Module or the old Embedding Module if new_num_tokens is None
        """
        if pad_to_multiple_of is not None:
            if not isinstance(pad_to_multiple_of, int):
                raise ValueError(
                    f"Asking to pad the embedding matrix to a multiple of `{pad_to_multiple_of}`, which is not and integer. Please make sure to pass an integer"
                )
            if new_num_tokens is None:
                new_num_tokens = old_embeddings.weight.shape[0]
            new_num_tokens = ((new_num_tokens + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
        else:
            logger.info(
                "You are resizing the embedding layer without providing a `pad_to_multiple_of` parameter. This means that the new embedding"
                f" dimension will be {new_num_tokens}. This might induce some performance reduction as *Tensor Cores* will not be available."
                " For more details about this, or help on choosing the correct value for resizing, refer to this guide:"
                " https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc"
            )

        if new_num_tokens is None:
            return old_embeddings

        old_num_tokens, old_embedding_dim = old_embeddings.weight.shape
        if old_num_tokens == new_num_tokens:
            return old_embeddings

        # Build new embeddings
        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)

        # initialize all new embeddings (in particular added tokens)
        self._init_weights(new_embeddings)

        # Copy word embeddings from the previous weights
        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)
        new_embeddings.weight.data[:num_tokens_to_copy, :] = old_embeddings.weight.data[
                                                                      :num_tokens_to_copy, :]
        return new_embeddings

    def _get_resized_lm_head(
        self, old_lm_head: nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False
    ) -> nn.Linear:
        """
        Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized
        vectors at the end. Reducing the size will remove vectors from the end

        Args:
            old_lm_head (`nn.Linear`):
                Old lm head liner layer to be resized.
            new_num_tokens (`int`, *optional*):
                New number of tokens in the linear 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
                `nn.Linear` module of the model without doing anything. transposed (`bool`, *optional*, defaults
                to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is `lm_head_dim,
                vocab_size` else `vocab_size, lm_head_dim`.

        Returns:
            `nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is
                `None`
        """
        if new_num_tokens is None:
            return old_lm_head

        old_num_tokens, old_lm_head_dim = (
            old_lm_head.weight.shape if not transposed else old_lm_head.weight.T.shape
        )

        if old_num_tokens == new_num_tokens:
            return old_lm_head

        if not isinstance(old_lm_head, nn.Linear):
            raise TypeError(
                f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. You"
                " should either use a different resize function or make sure that `old_lm_head` are an instance of"
                f" {nn.Linear}."
            )

        # Build new lm head
        new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim)
        has_new_lm_head_bias = old_lm_head.bias is not None

        # When using DeepSpeed ZeRO-3, we shouldn't create new embeddings with DeepSpeed init
        # because the shape of the new embedding layer is used across various modeling files
        # as well as to update config vocab size. Shape will be 0 when using DeepSpeed init leading
        # to errors when training.
        new_lm_head = nn.Linear(
            *new_lm_head_shape,
            bias=has_new_lm_head_bias,
        )

        # initialize new lm head (in particular added tokens)
        self._init_weights(new_lm_head)

        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)

        self._copy_lm_head_original_to_resized(
            new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias
        )

        return new_lm_head

    def _copy_lm_head_original_to_resized(
        self, new_lm_head, old_lm_head, num_tokens_to_copy, transposed, has_new_lm_head_bias
    ):
        """
        Copies the original language model head to a resized language model head.

        Args:
            self (PreTrainedModel): The instance of the PreTrainedModel class.
            new_lm_head (torch.nn.Module): The resized language model head to be copied into.
            old_lm_head (torch.nn.Module): The original language model head to be copied from.
            num_tokens_to_copy (int): The number of tokens to copy from the original language model head.
            transposed (bool): Whether the weight tensor of the new language model head is transposed.
            has_new_lm_head_bias (bool): Whether the new language model head has a bias tensor.

        Returns:
            None.

        Raises:
            None.
        """
        # Copy old lm head weights to new lm head
        if not transposed:
            new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
        else:
            new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]

        # Copy bias weights to new lm head
        if has_new_lm_head_bias:
            new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]

    @classmethod
    def load(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
             *args, **kwargs):
        """
        Load a pre-trained checkpoint from a pre-trained model file or url,
        download and cache the pre-trained model file if model name in model list.

        Params:
            pretrained_model_name_or_path:
        """
        return cls.from_pretrained(pretrained_model_name_or_path, args, kwargs)

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
        *model_args,
        config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
        cache_dir: Optional[Union[str, os.PathLike]] = None,
        ignore_mismatched_sizes: bool = False,
        force_download: bool = False,
        local_files_only: bool = False,
        token: Optional[Union[str, bool]] = None,
        use_safetensors: bool = None,
        mirror: str = 'huggingface',
        **kwargs,
    ):
        """from_pretrained"""
        state_dict = kwargs.pop("state_dict", None)
        cache_dir = kwargs.pop("cache_dir", None)
        _ = kwargs.pop("from_pt", True)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        _fast_init = kwargs.pop("_fast_init", True)
        output_loading_info = kwargs.pop("output_loading_info", False)
        subfolder = kwargs.pop("subfolder", "")
        variant = kwargs.pop("variant", None)
        ms_dtype = kwargs.pop("ms_dtype", None)
        _ = kwargs.pop('low_cpu_mem_usage', None)
        revision = kwargs.pop('revision', 'main')

        if use_safetensors is None and not is_safetensors_available():
            use_safetensors = False

        is_sharded = False
        # Load config if we don't provide a configuration
        if not isinstance(config, PretrainedConfig):
            config_path = config if config is not None else pretrained_model_name_or_path
            config, model_kwargs = cls.config_class.from_pretrained(
                config_path,
                *model_args,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                mirror=mirror,
                **kwargs,
            )
        else:
            model_kwargs = kwargs

        # Load model
        if pretrained_model_name_or_path is not None:
            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
            is_local = os.path.isdir(pretrained_model_name_or_path)
            if is_local:
                if os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, PT_WEIGHTS_NAME)
                ):
                    # Load from a PyTorch checkpoint
                    archive_file = os.path.join(pretrained_model_name_or_path, subfolder, PT_WEIGHTS_NAME)
                elif os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant))
                ):
                    # Load from a MindSpore checkpoint
                    archive_file = os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)
                    )
                elif use_safetensors is not False and os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant))
                ):
                    # Load from a safetensors checkpoint
                    archive_file = os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)
                    )
                elif os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(PT_WEIGHTS_INDEX_NAME, variant))
                ):
                    # Load from a sharded PyTorch checkpoint
                    archive_file = os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(PT_WEIGHTS_INDEX_NAME, variant)
                    )
                    is_sharded = True
                elif os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant))
                ):
                    # Load from a sharded MindSpore checkpoint
                    archive_file = os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)
                    )
                    is_sharded = True
                elif use_safetensors is not False and os.path.isfile(
                    os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)
                    )
                ):
                    # Load from a sharded safetensors checkpoint
                    archive_file = os.path.join(
                        pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)
                    )
                    is_sharded = True
                # At this stage we don't have a weight file so we will raise an error.
                elif use_safetensors:
                    raise EnvironmentError(
                        f"Error no file named {_add_variant(SAFE_WEIGHTS_NAME, variant)} found in directory"
                        f" {pretrained_model_name_or_path}."
                    )
                else:
                    raise EnvironmentError(
                        f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {PT_WEIGHTS_NAME},"
                        f" found in directory {pretrained_model_name_or_path}."
                    )
            elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
                archive_file = pretrained_model_name_or_path
                is_local = True
            elif is_remote_url(pretrained_model_name_or_path):
                filename = pretrained_model_name_or_path
                resolved_archive_file = download_url(pretrained_model_name_or_path)
            else:
                if use_safetensors is not False:
                    filename = _add_variant(SAFE_WEIGHTS_NAME, variant)
                else:
                    filename = _add_variant(WEIGHTS_NAME, variant)

                try:
                    # Load from URL or cache if already cached
                    cached_file_kwargs = {
                        "cache_dir": cache_dir,
                        "force_download": force_download,
                        "proxies": proxies,
                        "resume_download": resume_download,
                        "local_files_only": local_files_only,
                        "subfolder": subfolder,
                        "_raise_exceptions_for_missing_entries": False,
                        'revision': revision,
                        "token": token,
                        'mirror': mirror
                    }
                    # try safetensors
                    resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)
                    use_safetensors = resolved_archive_file is not None

                    # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None
                    # result when internet is up, the repo and revision exist, but the file does not.
                    if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):
                        # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path,
                            _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),
                            **cached_file_kwargs,
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True
                            use_safetensors = True

                    if resolved_archive_file is None:
                        filename = _add_variant(WEIGHTS_NAME, variant)
                        resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                    if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):
                        # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path,
                            _add_variant(WEIGHTS_INDEX_NAME, variant),
                            **cached_file_kwargs,
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    if resolved_archive_file is None:
                        filename = _add_variant(PT_WEIGHTS_NAME, variant)
                        resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                    if resolved_archive_file is None and filename == _add_variant(PT_WEIGHTS_NAME, variant):
                        # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                        resolved_archive_file = cached_file(
                            pretrained_model_name_or_path,
                            _add_variant(PT_WEIGHTS_INDEX_NAME, variant),
                            **cached_file_kwargs,
                        )
                        if resolved_archive_file is not None:
                            is_sharded = True

                    if resolved_archive_file is None:
                        raise EnvironmentError(
                            f"{pretrained_model_name_or_path} does not appear to have a file named"
                            f" {_add_variant(SAFE_WEIGHTS_NAME, variant)}, {_add_variant(PT_WEIGHTS_NAME, variant)}"
                        )
                except EnvironmentError:
                    # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
                    # to the original exception.
                    raise
                except Exception as exc:
                    # For any other exception, we throw a generic error.
                    raise EnvironmentError(
                        f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
                        ", make sure you don't have a local directory with the"
                        f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                        f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)},"
                        f" {_add_variant(PT_WEIGHTS_NAME, variant)}."
                    ) from exc

            if is_local:
                logger.info(f"loading weights file {archive_file}")
                resolved_archive_file = archive_file
            else:
                logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")
        else:
            resolved_archive_file = None

        if is_sharded:
            # rsolved_archive_file becomes a list of files that point to the different checkpoint shards in this case.
            resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
                pretrained_model_name_or_path,
                resolved_archive_file,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                resume_download=resume_download,
                local_files_only=local_files_only,
                token=token,
                subfolder=subfolder,
                revision=revision,
                mirror=mirror,
            )

        if pretrained_model_name_or_path is None and state_dict is None:
            raise ValueError("the argument 'pretrained_model_name_or_path' should be "
                             "a string of model name or checkpoint path, but got 'None'.")

        config.name_or_path = pretrained_model_name_or_path
        # Instantiate model.

        config_dict = config.to_dict()

        dtype_group = {key: getattr(config, key).ms_dtype for key in config_dict.keys() \
                       if isinstance(config_dict[key], dict) and 'ms_dtype' in config_dict[key]}

        if ms_dtype is None or ms_dtype == 'auto':
            ms_dtype = config.ms_dtype

        if ms_dtype is None:
            ms_dtype = mindspore.float32

        use_fp16 = False
        usage_dtype = mindspore.dtype_to_nptype(ms_dtype)
        if ms_dtype == mindspore.bfloat16:
            ms_dtype = mindspore.float16
            usage_dtype = np.float16
            use_fp16 = True

        def empty_initializer(init, shape=None, dtype=mindspore.float32):
            if not isinstance(shape, (tuple, list)):
                shape = (shape,)
            if dtype in (mindspore.float16, mindspore.float32) \
                and ms_dtype is not None:
                dtype = ms_dtype
            return Tensor_(shape=shape, dtype=dtype)

        with no_init_weights(empty_initializer, _fast_init):
            model = cls(config, *model_args, **model_kwargs)

        if is_sharded:
            converted_filenames = resolved_archive_file

        # tie the model weights before retrieving the state_dict
        model.tie_weights()

        ptrs = collections.defaultdict(list)
        for name, tensor in model.parameters_dict().items():
            id_tensor = hash(tensor)
            ptrs[id_tensor].append(name)

        # These are all the pointers of shared tensors.
        tied_params = [names for _, names in ptrs.items() if len(names) > 1]
        def load_ckpt(resolved_archive_file):
            if not resolved_archive_file.endswith('ckpt'):
                if use_safetensors or 'safetensors' in resolved_archive_file:
                    from safetensors.numpy import load_file
                    origin_state_dict = load_file(resolved_archive_file)
                    if use_fp16:
                        logger.warning_once("MindSpore do not support bfloat16 dtype, we will automaticlly convert to float16")
                    state_dict = {k: Parameter(Tensor.from_numpy(v.astype(usage_dtype))) for k, v in origin_state_dict.items()}
                else:
                    state_dict = load(resolved_archive_file)
            else:
                try:
                    state_dict = load_checkpoint(str(resolved_archive_file))
                except Exception as exc:
                    raise OSError(
                        f"Unable to load weights from mindspore checkpoint file '{resolved_archive_file}'. "
                    ) from exc

            state_keys = list(state_dict.keys())
            for key in state_keys:
                new_key = key.replace('gamma', 'weight').replace('beta', 'bias').replace('embedding_table', 'weight')
                if new_key != key:
                    state_dict[new_key] = state_dict.pop(key)
            return state_dict

        keys_missing = list(model.parameters_dict().keys())

        use_keep_in_fp32_modules = False
        if model._keep_in_fp32_modules:
            use_keep_in_fp32_modules = True

        remove_prefix_from_model = None
        add_prefix_to_model = None

        def fix_weight_norm_missing_keys(state_dict_keys: dict, keys_missing:List[str]) -> List[str]:
            ''' if both `weight_g` and `weight_v` are loaded, key `weight` is not missing :) '''
            non_missing_keys = []
            for key in keys_missing:
                if f'{key}_g' in state_dict_keys and f'{key}_v' in state_dict_keys:
                    non_missing_keys.append(key)
            return non_missing_keys

        def load_param_into_net(model: nn.Module, param_dict: dict, prefix: str, dtype_group: dict = None):
            state_dict_keys = list(param_dict.keys())
            keep_in_fp32_modules = model._keep_in_fp32_modules
            keys_unexpected = list(param_dict.keys())

            has_prefix_module = any(s.startswith(prefix) for s in keys_unexpected)
            expects_prefix_module = any(s.startswith(prefix) for s in keys_missing)

            nonlocal remove_prefix_from_model
            nonlocal add_prefix_to_model
            remove_prefix_from_model = not has_prefix_module and expects_prefix_module
            add_prefix_to_model = has_prefix_module and not expects_prefix_module

            for pname_in_net, param in model.state_dict().items():
                if add_prefix_to_model:
                    param_name = prefix + '.' + pname_in_net
                elif remove_prefix_from_model:
                    param_name = pname_in_net.replace(f'{prefix}.', '', 1)
                else:
                    param_name = pname_in_net

                new_param = param_dict.pop(param_name, None)
                module_dtype = None
                for m_name, m_dtype in dtype_group.items():
                    if m_name in param_name:
                        module_dtype = m_dtype
                        break

                if new_param is not None:
                    use_replace = False
                    if new_param.shape != param.shape:
                        if not ignore_mismatched_sizes:
                            raise RuntimeError(f'The shape of parameter `{param.name} is {param.shape}, but got mismatch parameter'
                                            f' `{param_name} with shape {new_param.shape} in checkpoint, '
                                            f'\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.')
                        logger.warning(f'The shape of parameter `{param.name} is {param.shape}, but got mismatch parameter'
                                        f' `{param_name} with shape {new_param.shape} in checkpoint, ')
                        continue

                    if use_keep_in_fp32_modules and \
                        any(module_to_keep_in_fp32 in pname_in_net.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules):
                        new_param = new_param.astype(mindspore.float32)
                    elif module_dtype and param.dtype in (mindspore.float32, mindspore.float16):
                        new_param = new_param.astype(module_dtype)
                    elif ms_dtype and param.dtype in (mindspore.float32, mindspore.float16):
                        new_param = new_param.astype(ms_dtype)

                    if new_param.dtype != param.dtype or new_param.shape != param.shape:
                        use_replace = True

                    if use_replace:
                        if isinstance(new_param, Parameter) and isinstance(param, Parameter):
                            new_param.name = param.name
                            new_param.requires_grad = param.requires_grad
                            replace_references(param, new_param)
                        elif isinstance(param, Tensor):
                            param.assign_value(new_param)
                        else:
                            replace_references(param, Parameter(new_param, requires_grad=param.requires_grad))
                    else:
                        param.assign_value(new_param)
                    keys_unexpected.remove(param_name)
                    if pname_in_net in keys_missing:
                        keys_missing.remove(pname_in_net)
                else:
                    # fix missing value parameter dtype cast.
                    if isinstance(param, Parameter) and ms_dtype and ms_dtype != param.dtype:
                        new_param = param.astype(ms_dtype)
                        replace_references(param, Parameter(new_param, requires_grad=param.requires_grad))

            # NOTE: monkey patching weight_norm
            for key in fix_weight_norm_missing_keys(state_dict_keys, keys_missing):
                keys_missing.remove(key)

            return keys_unexpected, keys_missing

        all_keys_unexpected = None
        if state_dict is None:
            if is_sharded:
                all_keys_unexpected = []
                for name in tqdm(converted_filenames, desc="Loading checkpoint shards"):
                    state_dict = load_ckpt(name)
                    keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)
                    all_keys_unexpected.extend(keys_unexpected)
                    del state_dict
                    gc.collect()
                loaded_keys = sharded_metadata["all_checkpoint_keys"]
            else:
                state_dict = load_ckpt(resolved_archive_file)
                loaded_keys = list(state_dict.keys())
                all_keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)
        else:
            loaded_keys = list(state_dict.keys())
            all_keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)

        loaded_add_keys = []
        for group in tied_params:
            missing_in_group = [k for k in keys_missing if k in group]
            if len(missing_in_group) > 0 and len(missing_in_group) < len(group):
                loaded_add_keys.extend([k for k in keys_missing if k in missing_in_group])
                keys_missing = [k for k in keys_missing if k not in missing_in_group]
        if cls._keys_to_ignore_on_load_missing is not None:
            for pat in cls._keys_to_ignore_on_load_missing:
                keys_missing = [k for k in keys_missing if re.search(pat, k) is None]

        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                all_keys_unexpected = [k for k in all_keys_unexpected if re.search(pat, k) is None]

        # make sure token embedding weights are still tied if needed
        model.tie_weights()

        # retrieve unintialized modules and initialize before maybe overriding that with the pretrained weights.
        if _fast_init:
            if not ignore_mismatched_sizes:
                if remove_prefix_from_model:
                    _loaded_keys = [f"{cls.base_model_prefix}.{k}" for k in loaded_keys]
                elif add_prefix_to_model:
                    _loaded_keys = [k[len(cls.base_model_prefix) + 1 :] for k in loaded_keys]
                else:
                    _loaded_keys = loaded_keys

                _loaded_keys += loaded_add_keys
                _ = set_initialized_submodules(model, _loaded_keys)
            else:
                _ = dict(model.cells_and_names())

            model.apply(model._initialize_weights)

        # Set model in evaluation mode to deactivate DropOut modules by default
        model.set_train(False)

        # If it is a model with generation capabilities, attempt to load the generation config
        if model.can_generate() and pretrained_model_name_or_path is not None:
            try:
                model.generation_config = GenerationConfig.from_pretrained(
                    pretrained_model_name_or_path,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    resume_download=resume_download,
                    proxies=proxies,
                    local_files_only=local_files_only,
                    subfolder=subfolder,
                    revision=revision,
                    mirror=mirror,
                    **kwargs,
                )
            except OSError:
                logger.info(
                    "Generation config file not found, using a generation config created from the model config."
                )

        if output_loading_info:
            loading_info = {
                "missing_keys": keys_missing,
                "unexpected_keys": all_keys_unexpected,
            }
            return model, loading_info

        if all_keys_unexpected:
            logger.warning(f'The following parameters in checkpoint files are not loaded:\n'
                           f'{all_keys_unexpected}')
        if keys_missing:
            logger.warning(f'The following parameters in models are missing parameter:\n'
                           f'{keys_missing}')
        return model

    def save(self, save_dir):
        """
        Save a model and its configuration file to a directory, so that
        it can be re-loaded using the `:func:`PreTrainedModel.from_pretrained`` class method.

        Arguments:
            save_dir: directory to which to save.
        """
        if os.path.isfile(save_dir):
            logger.error(f"Provided path ({save_dir}) should be a directory, not a file")
            return
        os.makedirs(save_dir, exist_ok=True)

        # Only save the model itself if we are using distributed training
        model_to_save = self.cell if hasattr(self, "cell") else self

        # Attach architecture to the config
        model_to_save.config.architectures = [model_to_save.__class__.__name__]

        # If we save using the predefined names, we can load using `from_pretrained`
        output_model_file = os.path.join(save_dir, WEIGHTS_NAME)
        save_checkpoint(model_to_save, output_model_file)

        logger.info(f"Model weights saved in {output_model_file}")

    @classmethod
    def can_generate(cls) -> 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.
        # Alternativelly, the model can also have a custom `generate` function.
        if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate):
            return False
        return True

    def warn_if_padding_and_no_attention_mask(self, input_ids, attention_mask):
        """
        Shows a one-time warning if the input_ids appear to contain padding and no attention mask was given.
        """
        if (attention_mask is not None) or (self.config.pad_token_id is None):
            return

        # Check only the first and last input IDs to reduce overhead.
        # if self.config.pad_token_id in input_ids[:, [-1, 0]]:
        if ops.contains(input_ids[:, [-1, 0]], self.config.pad_token_id):
            warn_string = (
                "We strongly recommend passing in an `attention_mask` since your input_ids may be padded."
            )

            # If the pad token is equal to either BOS, EOS, or SEP, we do not know whether the user should use an
            # attention_mask or not. In this case, we should still show a warning because this is a rare case.
            if (
                (self.config.bos_token_id is not None and self.config.bos_token_id == self.config.pad_token_id)
                or (self.config.eos_token_id is not None and self.config.eos_token_id == self.config.pad_token_id)
                or (self.config.sep_token_id is not None and self.config.sep_token_id == self.config.pad_token_id)
            ):
                warn_string += (
                    f"\nYou may ignore this warning if your `pad_token_id` ({self.config.pad_token_id}) is identical "
                    f"to the `bos_token_id` ({self.config.bos_token_id}), `eos_token_id` ({self.config.eos_token_id}), "
                    f"or the `sep_token_id` ({self.config.sep_token_id}), and your input is not padded."
                )

            logger.warning(warn_string)

    def num_parameters(self, only_trainable=False):
        """return parameters count"""
        total = 0
        for param in self.get_parameters():
            if (only_trainable or param.requires_grad):
                total += param.size
        return total

    def trainable_params(self, recurse=True):
        """
        fix duplicated weights
        """
        return list(set(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse))))

    def check_names_and_refresh_name(self):
        """
        fix ignore tied weights
        """
        if not hasattr(self, "_params"):
            return
        all_name = dict(self.parameters_and_names()).keys()

        if len(set(all_name)) < len(all_name):
            self.update_parameters_name()
            self.check_names()

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        is_main_process: bool = True,
        state_dict: Optional[dict] = None,
        save_function: Callable = mindspore.save_checkpoint,
        max_shard_size: Union[int, str] = "5GB",
        safe_serialization: bool = True,
        variant: Optional[str] = None,
        **kwargs,
    ):
        """
        Save a model and its configuration file to a directory, so that it can be re-loaded using the
        [`~PreTrainedModel.from_pretrained`] class method.

        Arguments:
            save_directory (`str` or `os.PathLike`):
                Directory to which to save. Will be created if it doesn't exist.
            is_main_process (`bool`, *optional*, defaults to `True`):
                Whether the process calling this is the main process or not. Useful when in distributed training like
                TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
                the main process to avoid race conditions.
            state_dict (nested dictionary of `torch.Tensor`):
                The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only
                save parts of the model or if special precautions need to be taken when recovering the state dictionary
                of a model (like when using model parallelism).
            save_function (`Callable`):
                The function to use to save the state dictionary. Useful on distributed training like TPUs when one
                need to replace `torch.save` by another method.
            max_shard_size (`int` or `str`, *optional*, defaults to `"5GB"`):
                The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
                lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
                We default it to 5GB in order for models to be able to run easily on free-tier google colab instances
                without CPU OOM issues.

                <Tip warning={true}>

                If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
                which will be bigger than `max_shard_size`.

                </Tip>
            variant (`str`, *optional*):
                If specified, weights are saved in the format pytorch_model.<variant>.bin.
            save_peft_format (`bool`, *optional*, defaults to `True`):
                For backward compatibility with PEFT library, in case adapter weights are attached to the model, all
                keys of the state dict of adapters needs to be pre-pended with `base_model.model`. Advanced users can
                disable this behaviours by setting `save_peft_format` to `False`.
            kwargs (`Dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """
        if os.path.isfile(save_directory):
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
            return

        os.makedirs(save_directory, exist_ok=True)

        # Only save the model itself if we are using distributed training
        model_to_save = self

        # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
        # we currently don't use this setting automatically, but may start to use with v5
        dtype = get_parameter_dtype(model_to_save)
        model_to_save.config.ms_dtype = str(dtype).lower()

        # Attach architecture to the config
        model_to_save.config.architectures = [model_to_save.__class__.__name__]

        # Save the config
        if is_main_process:
            model_to_save.config.save_pretrained(save_directory)
            if self.can_generate():
                model_to_save.generation_config.save_pretrained(save_directory)

        # Save the model
        if state_dict is None:
            state_dict = model_to_save.parameters_dict()

        # Handle the case where some state_dict keys shouldn't be saved
        if self._keys_to_ignore_on_save is not None:
            for ignore_key in self._keys_to_ignore_on_save:
                if ignore_key in state_dict.keys():
                    del state_dict[ignore_key]

        # Shard the model if it is too big.
        # weights_name = _add_variant(WEIGHTS_NAME, variant)
        weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME

        shards, index = shard_checkpoint(state_dict, max_shard_size=max_shard_size, weights_name=weights_name)

        # Clean the folder from a previous save
        for filename in os.listdir(save_directory):
            full_filename = os.path.join(save_directory, filename)
            # If we have a shard file that is not going to be replaced, we delete it, but only from the main process
            # in distributed settings to avoid race conditions.
            weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")

            # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
            filename_no_suffix = filename.replace(".bin", "").replace(".safetensors", "")
            reg = re.compile(r"(.*?)-\d{5}-of-\d{5}")

            if (
                filename.startswith(weights_no_suffix)
                and os.path.isfile(full_filename)
                and filename not in shards
                and is_main_process
                and reg.fullmatch(filename_no_suffix) is not None
            ):
                os.remove(full_filename)

        # Save the model
        for shard_file, shard in shards.items():
            if safe_serialization:
                # At some point we will need to deal better with save_function (used for TPU and other distributed
                # joyfulness), but for now this enough.
                safe_save_file(shard, os.path.join(save_directory, shard_file), metadata={"format": "np"})
            else:
                save_function(shard, os.path.join(save_directory, shard_file))

        if index is None:
            path_to_weights = os.path.join(save_directory, _add_variant(WEIGHTS_NAME, variant))
            logger.info(f"Model weights saved in {path_to_weights}")
        else:
            save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
            save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
            # Save the index as well
            with open(save_index_file, "w", encoding="utf-8") as f:
                content = json.dumps(index, indent=2, sort_keys=True) + "\n"
                f.write(content)
            logger.info(
                f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
                f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
                f"index located at {save_index_file}."
            )

    def enable_recompute(self):
        """Activates recompute (aka gradient checkpointing) for the current model."""
        if not self.supports_recompute:
            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")

        for _, cell in self.cells_and_names():
            if hasattr(cell, "_set_recompute"):
                cell._set_recompute()

    def check_names(self):
        """
        This method checks the names in the PreTrainedModel class.

        Args:
            self: The instance of the PreTrainedModel class.

        Returns:
            None.

        Raises:
            None.
        """

mindnlp.transformers.modeling_utils.PreTrainedModel.base_model property

to get base_model

mindnlp.transformers.modeling_utils.PreTrainedModel.__init__(config)

Initializes a new instance of the PreTrainedModel class.

PARAMETER	DESCRIPTION
self	The instance of the class.
config	A dictionary containing the configuration parameters for the model.

RETURNS	DESCRIPTION
	None.

RAISES	DESCRIPTION
ValueError	If the config parameter is invalid or missing required fields.
TypeError	If the config parameter is not of the expected type.
RuntimeError	If an error occurs during the initialization process.

Source code in mindnlp/transformers/modeling_utils.py

def __init__(self, config):
    """

    Initializes a new instance of the PreTrainedModel class.

    Args:
        self: The instance of the class.
        config: A dictionary containing the configuration parameters for the model.

    Returns:
        None.

    Raises:
        ValueError: If the config parameter is invalid or missing required fields.
        TypeError: If the config parameter is not of the expected type.
        RuntimeError: If an error occurs during the initialization process.
    """
    super().__init__()
    self._check_and_unset_acl()
    # Save config in model
    self.config = config
    self.name_or_path = config.name_or_path
    self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None

mindnlp.transformers.modeling_utils.PreTrainedModel.can_generate() classmethod

Returns whether this model can generate sequences with .generate().

RETURNS	DESCRIPTION
bool	bool: Whether this model can generate sequences with .generate().

Source code in mindnlp/transformers/modeling_utils.py

@classmethod
def can_generate(cls) -> 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.
    # Alternativelly, the model can also have a custom `generate` function.
    if "GenerationMixin" in str(cls.prepare_inputs_for_generation) and "GenerationMixin" in str(cls.generate):
        return False
    return True

mindnlp.transformers.modeling_utils.PreTrainedModel.check_names()

This method checks the names in the PreTrainedModel class.

PARAMETER	DESCRIPTION
self	The instance of the PreTrainedModel class.

RETURNS	DESCRIPTION
	None.

Source code in mindnlp/transformers/modeling_utils.py

def check_names(self):
    """
    This method checks the names in the PreTrainedModel class.

    Args:
        self: The instance of the PreTrainedModel class.

    Returns:
        None.

    Raises:
        None.
    """

mindnlp.transformers.modeling_utils.PreTrainedModel.check_names_and_refresh_name()

fix ignore tied weights

Source code in mindnlp/transformers/modeling_utils.py

def check_names_and_refresh_name(self):
    """
    fix ignore tied weights
    """
    if not hasattr(self, "_params"):
        return
    all_name = dict(self.parameters_and_names()).keys()

    if len(set(all_name)) < len(all_name):
        self.update_parameters_name()
        self.check_names()

mindnlp.transformers.modeling_utils.PreTrainedModel.enable_recompute()

Activates recompute (aka gradient checkpointing) for the current model.

Source code in mindnlp/transformers/modeling_utils.py

def enable_recompute(self):
    """Activates recompute (aka gradient checkpointing) for the current model."""
    if not self.supports_recompute:
        raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")

    for _, cell in self.cells_and_names():
        if hasattr(cell, "_set_recompute"):
            cell._set_recompute()

mindnlp.transformers.modeling_utils.PreTrainedModel.from_pretrained(pretrained_model_name_or_path, *model_args, config=None, cache_dir=None, ignore_mismatched_sizes=False, force_download=False, local_files_only=False, token=None, use_safetensors=None, mirror='huggingface', **kwargs) classmethod

from_pretrained

Source code in mindnlp/transformers/modeling_utils.py

@classmethod
def from_pretrained(
    cls,
    pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
    *model_args,
    config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,
    cache_dir: Optional[Union[str, os.PathLike]] = None,
    ignore_mismatched_sizes: bool = False,
    force_download: bool = False,
    local_files_only: bool = False,
    token: Optional[Union[str, bool]] = None,
    use_safetensors: bool = None,
    mirror: str = 'huggingface',
    **kwargs,
):
    """from_pretrained"""
    state_dict = kwargs.pop("state_dict", None)
    cache_dir = kwargs.pop("cache_dir", None)
    _ = kwargs.pop("from_pt", True)
    force_download = kwargs.pop("force_download", False)
    resume_download = kwargs.pop("resume_download", False)
    proxies = kwargs.pop("proxies", None)
    local_files_only = kwargs.pop("local_files_only", False)
    _fast_init = kwargs.pop("_fast_init", True)
    output_loading_info = kwargs.pop("output_loading_info", False)
    subfolder = kwargs.pop("subfolder", "")
    variant = kwargs.pop("variant", None)
    ms_dtype = kwargs.pop("ms_dtype", None)
    _ = kwargs.pop('low_cpu_mem_usage', None)
    revision = kwargs.pop('revision', 'main')

    if use_safetensors is None and not is_safetensors_available():
        use_safetensors = False

    is_sharded = False
    # Load config if we don't provide a configuration
    if not isinstance(config, PretrainedConfig):
        config_path = config if config is not None else pretrained_model_name_or_path
        config, model_kwargs = cls.config_class.from_pretrained(
            config_path,
            *model_args,
            cache_dir=cache_dir,
            return_unused_kwargs=True,
            force_download=force_download,
            resume_download=resume_download,
            proxies=proxies,
            local_files_only=local_files_only,
            mirror=mirror,
            **kwargs,
        )
    else:
        model_kwargs = kwargs

    # Load model
    if pretrained_model_name_or_path is not None:
        pretrained_model_name_or_path = str(pretrained_model_name_or_path)
        is_local = os.path.isdir(pretrained_model_name_or_path)
        if is_local:
            if os.path.isfile(
                os.path.join(pretrained_model_name_or_path, subfolder, PT_WEIGHTS_NAME)
            ):
                # Load from a PyTorch checkpoint
                archive_file = os.path.join(pretrained_model_name_or_path, subfolder, PT_WEIGHTS_NAME)
            elif os.path.isfile(
                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant))
            ):
                # Load from a MindSpore checkpoint
                archive_file = os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_NAME, variant)
                )
            elif use_safetensors is not False and os.path.isfile(
                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant))
            ):
                # Load from a safetensors checkpoint
                archive_file = os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_NAME, variant)
                )
            elif os.path.isfile(
                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(PT_WEIGHTS_INDEX_NAME, variant))
            ):
                # Load from a sharded PyTorch checkpoint
                archive_file = os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(PT_WEIGHTS_INDEX_NAME, variant)
                )
                is_sharded = True
            elif os.path.isfile(
                os.path.join(pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant))
            ):
                # Load from a sharded MindSpore checkpoint
                archive_file = os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(WEIGHTS_INDEX_NAME, variant)
                )
                is_sharded = True
            elif use_safetensors is not False and os.path.isfile(
                os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)
                )
            ):
                # Load from a sharded safetensors checkpoint
                archive_file = os.path.join(
                    pretrained_model_name_or_path, subfolder, _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant)
                )
                is_sharded = True
            # At this stage we don't have a weight file so we will raise an error.
            elif use_safetensors:
                raise EnvironmentError(
                    f"Error no file named {_add_variant(SAFE_WEIGHTS_NAME, variant)} found in directory"
                    f" {pretrained_model_name_or_path}."
                )
            else:
                raise EnvironmentError(
                    f"Error no file named {_add_variant(WEIGHTS_NAME, variant)}, {PT_WEIGHTS_NAME},"
                    f" found in directory {pretrained_model_name_or_path}."
                )
        elif os.path.isfile(os.path.join(subfolder, pretrained_model_name_or_path)):
            archive_file = pretrained_model_name_or_path
            is_local = True
        elif is_remote_url(pretrained_model_name_or_path):
            filename = pretrained_model_name_or_path
            resolved_archive_file = download_url(pretrained_model_name_or_path)
        else:
            if use_safetensors is not False:
                filename = _add_variant(SAFE_WEIGHTS_NAME, variant)
            else:
                filename = _add_variant(WEIGHTS_NAME, variant)

            try:
                # Load from URL or cache if already cached
                cached_file_kwargs = {
                    "cache_dir": cache_dir,
                    "force_download": force_download,
                    "proxies": proxies,
                    "resume_download": resume_download,
                    "local_files_only": local_files_only,
                    "subfolder": subfolder,
                    "_raise_exceptions_for_missing_entries": False,
                    'revision': revision,
                    "token": token,
                    'mirror': mirror
                }
                # try safetensors
                resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)
                use_safetensors = resolved_archive_file is not None

                # Since we set _raise_exceptions_for_missing_entries=False, we don't get an exception but a None
                # result when internet is up, the repo and revision exist, but the file does not.
                if resolved_archive_file is None and filename == _add_variant(SAFE_WEIGHTS_NAME, variant):
                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                    resolved_archive_file = cached_file(
                        pretrained_model_name_or_path,
                        _add_variant(SAFE_WEIGHTS_INDEX_NAME, variant),
                        **cached_file_kwargs,
                    )
                    if resolved_archive_file is not None:
                        is_sharded = True
                        use_safetensors = True

                if resolved_archive_file is None:
                    filename = _add_variant(WEIGHTS_NAME, variant)
                    resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                if resolved_archive_file is None and filename == _add_variant(WEIGHTS_NAME, variant):
                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                    resolved_archive_file = cached_file(
                        pretrained_model_name_or_path,
                        _add_variant(WEIGHTS_INDEX_NAME, variant),
                        **cached_file_kwargs,
                    )
                    if resolved_archive_file is not None:
                        is_sharded = True

                if resolved_archive_file is None:
                    filename = _add_variant(PT_WEIGHTS_NAME, variant)
                    resolved_archive_file = cached_file(pretrained_model_name_or_path, filename, **cached_file_kwargs)

                if resolved_archive_file is None and filename == _add_variant(PT_WEIGHTS_NAME, variant):
                    # Maybe the checkpoint is sharded, we try to grab the index name in this case.
                    resolved_archive_file = cached_file(
                        pretrained_model_name_or_path,
                        _add_variant(PT_WEIGHTS_INDEX_NAME, variant),
                        **cached_file_kwargs,
                    )
                    if resolved_archive_file is not None:
                        is_sharded = True

                if resolved_archive_file is None:
                    raise EnvironmentError(
                        f"{pretrained_model_name_or_path} does not appear to have a file named"
                        f" {_add_variant(SAFE_WEIGHTS_NAME, variant)}, {_add_variant(PT_WEIGHTS_NAME, variant)}"
                    )
            except EnvironmentError:
                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted
                # to the original exception.
                raise
            except Exception as exc:
                # For any other exception, we throw a generic error.
                raise EnvironmentError(
                    f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it"
                    ", make sure you don't have a local directory with the"
                    f" same name. Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a"
                    f" directory containing a file named {_add_variant(WEIGHTS_NAME, variant)}, {_add_variant(SAFE_WEIGHTS_NAME, variant)},"
                    f" {_add_variant(PT_WEIGHTS_NAME, variant)}."
                ) from exc

        if is_local:
            logger.info(f"loading weights file {archive_file}")
            resolved_archive_file = archive_file
        else:
            logger.info(f"loading weights file {filename} from cache at {resolved_archive_file}")
    else:
        resolved_archive_file = None

    if is_sharded:
        # rsolved_archive_file becomes a list of files that point to the different checkpoint shards in this case.
        resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
            pretrained_model_name_or_path,
            resolved_archive_file,
            cache_dir=cache_dir,
            force_download=force_download,
            proxies=proxies,
            resume_download=resume_download,
            local_files_only=local_files_only,
            token=token,
            subfolder=subfolder,
            revision=revision,
            mirror=mirror,
        )

    if pretrained_model_name_or_path is None and state_dict is None:
        raise ValueError("the argument 'pretrained_model_name_or_path' should be "
                         "a string of model name or checkpoint path, but got 'None'.")

    config.name_or_path = pretrained_model_name_or_path
    # Instantiate model.

    config_dict = config.to_dict()

    dtype_group = {key: getattr(config, key).ms_dtype for key in config_dict.keys() \
                   if isinstance(config_dict[key], dict) and 'ms_dtype' in config_dict[key]}

    if ms_dtype is None or ms_dtype == 'auto':
        ms_dtype = config.ms_dtype

    if ms_dtype is None:
        ms_dtype = mindspore.float32

    use_fp16 = False
    usage_dtype = mindspore.dtype_to_nptype(ms_dtype)
    if ms_dtype == mindspore.bfloat16:
        ms_dtype = mindspore.float16
        usage_dtype = np.float16
        use_fp16 = True

    def empty_initializer(init, shape=None, dtype=mindspore.float32):
        if not isinstance(shape, (tuple, list)):
            shape = (shape,)
        if dtype in (mindspore.float16, mindspore.float32) \
            and ms_dtype is not None:
            dtype = ms_dtype
        return Tensor_(shape=shape, dtype=dtype)

    with no_init_weights(empty_initializer, _fast_init):
        model = cls(config, *model_args, **model_kwargs)

    if is_sharded:
        converted_filenames = resolved_archive_file

    # tie the model weights before retrieving the state_dict
    model.tie_weights()

    ptrs = collections.defaultdict(list)
    for name, tensor in model.parameters_dict().items():
        id_tensor = hash(tensor)
        ptrs[id_tensor].append(name)

    # These are all the pointers of shared tensors.
    tied_params = [names for _, names in ptrs.items() if len(names) > 1]
    def load_ckpt(resolved_archive_file):
        if not resolved_archive_file.endswith('ckpt'):
            if use_safetensors or 'safetensors' in resolved_archive_file:
                from safetensors.numpy import load_file
                origin_state_dict = load_file(resolved_archive_file)
                if use_fp16:
                    logger.warning_once("MindSpore do not support bfloat16 dtype, we will automaticlly convert to float16")
                state_dict = {k: Parameter(Tensor.from_numpy(v.astype(usage_dtype))) for k, v in origin_state_dict.items()}
            else:
                state_dict = load(resolved_archive_file)
        else:
            try:
                state_dict = load_checkpoint(str(resolved_archive_file))
            except Exception as exc:
                raise OSError(
                    f"Unable to load weights from mindspore checkpoint file '{resolved_archive_file}'. "
                ) from exc

        state_keys = list(state_dict.keys())
        for key in state_keys:
            new_key = key.replace('gamma', 'weight').replace('beta', 'bias').replace('embedding_table', 'weight')
            if new_key != key:
                state_dict[new_key] = state_dict.pop(key)
        return state_dict

    keys_missing = list(model.parameters_dict().keys())

    use_keep_in_fp32_modules = False
    if model._keep_in_fp32_modules:
        use_keep_in_fp32_modules = True

    remove_prefix_from_model = None
    add_prefix_to_model = None

    def fix_weight_norm_missing_keys(state_dict_keys: dict, keys_missing:List[str]) -> List[str]:
        ''' if both `weight_g` and `weight_v` are loaded, key `weight` is not missing :) '''
        non_missing_keys = []
        for key in keys_missing:
            if f'{key}_g' in state_dict_keys and f'{key}_v' in state_dict_keys:
                non_missing_keys.append(key)
        return non_missing_keys

    def load_param_into_net(model: nn.Module, param_dict: dict, prefix: str, dtype_group: dict = None):
        state_dict_keys = list(param_dict.keys())
        keep_in_fp32_modules = model._keep_in_fp32_modules
        keys_unexpected = list(param_dict.keys())

        has_prefix_module = any(s.startswith(prefix) for s in keys_unexpected)
        expects_prefix_module = any(s.startswith(prefix) for s in keys_missing)

        nonlocal remove_prefix_from_model
        nonlocal add_prefix_to_model
        remove_prefix_from_model = not has_prefix_module and expects_prefix_module
        add_prefix_to_model = has_prefix_module and not expects_prefix_module

        for pname_in_net, param in model.state_dict().items():
            if add_prefix_to_model:
                param_name = prefix + '.' + pname_in_net
            elif remove_prefix_from_model:
                param_name = pname_in_net.replace(f'{prefix}.', '', 1)
            else:
                param_name = pname_in_net

            new_param = param_dict.pop(param_name, None)
            module_dtype = None
            for m_name, m_dtype in dtype_group.items():
                if m_name in param_name:
                    module_dtype = m_dtype
                    break

            if new_param is not None:
                use_replace = False
                if new_param.shape != param.shape:
                    if not ignore_mismatched_sizes:
                        raise RuntimeError(f'The shape of parameter `{param.name} is {param.shape}, but got mismatch parameter'
                                        f' `{param_name} with shape {new_param.shape} in checkpoint, '
                                        f'\n\tYou may consider adding `ignore_mismatched_sizes=True` in the model `from_pretrained` method.')
                    logger.warning(f'The shape of parameter `{param.name} is {param.shape}, but got mismatch parameter'
                                    f' `{param_name} with shape {new_param.shape} in checkpoint, ')
                    continue

                if use_keep_in_fp32_modules and \
                    any(module_to_keep_in_fp32 in pname_in_net.split(".") for module_to_keep_in_fp32 in keep_in_fp32_modules):
                    new_param = new_param.astype(mindspore.float32)
                elif module_dtype and param.dtype in (mindspore.float32, mindspore.float16):
                    new_param = new_param.astype(module_dtype)
                elif ms_dtype and param.dtype in (mindspore.float32, mindspore.float16):
                    new_param = new_param.astype(ms_dtype)

                if new_param.dtype != param.dtype or new_param.shape != param.shape:
                    use_replace = True

                if use_replace:
                    if isinstance(new_param, Parameter) and isinstance(param, Parameter):
                        new_param.name = param.name
                        new_param.requires_grad = param.requires_grad
                        replace_references(param, new_param)
                    elif isinstance(param, Tensor):
                        param.assign_value(new_param)
                    else:
                        replace_references(param, Parameter(new_param, requires_grad=param.requires_grad))
                else:
                    param.assign_value(new_param)
                keys_unexpected.remove(param_name)
                if pname_in_net in keys_missing:
                    keys_missing.remove(pname_in_net)
            else:
                # fix missing value parameter dtype cast.
                if isinstance(param, Parameter) and ms_dtype and ms_dtype != param.dtype:
                    new_param = param.astype(ms_dtype)
                    replace_references(param, Parameter(new_param, requires_grad=param.requires_grad))

        # NOTE: monkey patching weight_norm
        for key in fix_weight_norm_missing_keys(state_dict_keys, keys_missing):
            keys_missing.remove(key)

        return keys_unexpected, keys_missing

    all_keys_unexpected = None
    if state_dict is None:
        if is_sharded:
            all_keys_unexpected = []
            for name in tqdm(converted_filenames, desc="Loading checkpoint shards"):
                state_dict = load_ckpt(name)
                keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)
                all_keys_unexpected.extend(keys_unexpected)
                del state_dict
                gc.collect()
            loaded_keys = sharded_metadata["all_checkpoint_keys"]
        else:
            state_dict = load_ckpt(resolved_archive_file)
            loaded_keys = list(state_dict.keys())
            all_keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)
    else:
        loaded_keys = list(state_dict.keys())
        all_keys_unexpected, keys_missing = load_param_into_net(model, state_dict, cls.base_model_prefix, dtype_group)

    loaded_add_keys = []
    for group in tied_params:
        missing_in_group = [k for k in keys_missing if k in group]
        if len(missing_in_group) > 0 and len(missing_in_group) < len(group):
            loaded_add_keys.extend([k for k in keys_missing if k in missing_in_group])
            keys_missing = [k for k in keys_missing if k not in missing_in_group]
    if cls._keys_to_ignore_on_load_missing is not None:
        for pat in cls._keys_to_ignore_on_load_missing:
            keys_missing = [k for k in keys_missing if re.search(pat, k) is None]

    if cls._keys_to_ignore_on_load_unexpected is not None:
        for pat in cls._keys_to_ignore_on_load_unexpected:
            all_keys_unexpected = [k for k in all_keys_unexpected if re.search(pat, k) is None]

    # make sure token embedding weights are still tied if needed
    model.tie_weights()

    # retrieve unintialized modules and initialize before maybe overriding that with the pretrained weights.
    if _fast_init:
        if not ignore_mismatched_sizes:
            if remove_prefix_from_model:
                _loaded_keys = [f"{cls.base_model_prefix}.{k}" for k in loaded_keys]
            elif add_prefix_to_model:
                _loaded_keys = [k[len(cls.base_model_prefix) + 1 :] for k in loaded_keys]
            else:
                _loaded_keys = loaded_keys

            _loaded_keys += loaded_add_keys
            _ = set_initialized_submodules(model, _loaded_keys)
        else:
            _ = dict(model.cells_and_names())

        model.apply(model._initialize_weights)

    # Set model in evaluation mode to deactivate DropOut modules by default
    model.set_train(False)

    # If it is a model with generation capabilities, attempt to load the generation config
    if model.can_generate() and pretrained_model_name_or_path is not None:
        try:
            model.generation_config = GenerationConfig.from_pretrained(
                pretrained_model_name_or_path,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                subfolder=subfolder,
                revision=revision,
                mirror=mirror,
                **kwargs,
            )
        except OSError:
            logger.info(
                "Generation config file not found, using a generation config created from the model config."
            )

    if output_loading_info:
        loading_info = {
            "missing_keys": keys_missing,
            "unexpected_keys": all_keys_unexpected,
        }
        return model, loading_info

    if all_keys_unexpected:
        logger.warning(f'The following parameters in checkpoint files are not loaded:\n'
                       f'{all_keys_unexpected}')
    if keys_missing:
        logger.warning(f'The following parameters in models are missing parameter:\n'
                       f'{keys_missing}')
    return model

mindnlp.transformers.modeling_utils.PreTrainedModel.get_input_embeddings()

Returns the model's input embeddings.

RETURNS	DESCRIPTION
Module

Source code in mindnlp/transformers/modeling_utils.py

def get_input_embeddings(self) -> "nn.Module":
    """
    Returns the model's input embeddings.

    Returns:
        :obj:`nn.Module`: A mindspore cell mapping vocabulary to hidden states.
    """
    base_model = getattr(self, self.base_model_prefix, self)
    print(base_model)
    if base_model is not self:
        return base_model.get_input_embeddings()
    raise NotImplementedError

mindnlp.transformers.modeling_utils.PreTrainedModel.get_output_embeddings()

Get model's output embeddings Return None if the model doesn't have output embeddings

Source code in mindnlp/transformers/modeling_utils.py

def get_output_embeddings(self):
    """ Get model's output embeddings
        Return None if the model doesn't have output embeddings
    """
    return None  # Overwrite for models with output embeddings

mindnlp.transformers.modeling_utils.PreTrainedModel.get_position_embeddings()

get the model position embeddings if necessary

Source code in mindnlp/transformers/modeling_utils.py

def get_position_embeddings(self):
    """
    get the model position embeddings if necessary
    """
    raise NotImplementedError(
        f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
        f"overwrite this method in the class {self.__class__}"
    )

mindnlp.transformers.modeling_utils.PreTrainedModel.init_weights()

If needed prunes and maybe initializes weights. If using a custom PreTrainedModel, you need to implement any initialization logic in _init_weights.

Source code in mindnlp/transformers/modeling_utils.py

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`.
    """
    # Prune heads if needed
    if self.config.pruned_heads:
        self.prune_heads(self.config.pruned_heads)

    if _init_weights:
        # Initialize weights
        if getattr(self, 'apply', None):
            self.apply(self._initialize_weights)
        else:
            for _, cell in self.name_cells().items():
                self._initialize_weights(cell)

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

mindnlp.transformers.modeling_utils.PreTrainedModel.load(pretrained_model_name_or_path, *args, **kwargs) classmethod

Load a pre-trained checkpoint from a pre-trained model file or url, download and cache the pre-trained model file if model name in model list.

PARAMETER	DESCRIPTION
pretrained_model_name_or_path	TYPE: Optional[Union[str, PathLike]]

Source code in mindnlp/transformers/modeling_utils.py

@classmethod
def load(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],
         *args, **kwargs):
    """
    Load a pre-trained checkpoint from a pre-trained model file or url,
    download and cache the pre-trained model file if model name in model list.

    Params:
        pretrained_model_name_or_path:
    """
    return cls.from_pretrained(pretrained_model_name_or_path, args, kwargs)

mindnlp.transformers.modeling_utils.PreTrainedModel.num_parameters(only_trainable=False)

return parameters count

Source code in mindnlp/transformers/modeling_utils.py

def num_parameters(self, only_trainable=False):
    """return parameters count"""
    total = 0
    for param in self.get_parameters():
        if (only_trainable or param.requires_grad):
            total += param.size
    return total

mindnlp.transformers.modeling_utils.PreTrainedModel.post_init()

A method executed at the end of each Transformer model initialization, to execute code that needs the model's modules properly initialized (such as weight initialization).

Source code in mindnlp/transformers/modeling_utils.py

def post_init(self):
    """
    A method executed at the end of each Transformer model initialization, to execute code that needs the model's
    modules properly initialized (such as weight initialization).

    """
    self.init_weights()

mindnlp.transformers.modeling_utils.PreTrainedModel.prune_heads(heads_to_prune)

Prunes heads of the base model.

PARAMETER	DESCRIPTION
heads_to_prune	Dictionary with keys being selected layer indices (int) and associated values being the list of heads to prune in said layer (list of int). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2. TYPE: `Dict[int, List[int]]`

Source code in mindnlp/transformers/modeling_utils.py

def prune_heads(self, heads_to_prune: Dict[int, List[int]]):
    """
    Prunes heads of the base model.

    Arguments:
        heads_to_prune (`Dict[int, List[int]]`):
            Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads
            to prune in said layer (list of `int`). For instance {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on
            layer 1 and heads 2 and 3 on layer 2.
    """
    # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
    for layer, heads in heads_to_prune.items():
        union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)
        self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSON

    self.base_model._prune_heads(heads_to_prune)

mindnlp.transformers.modeling_utils.PreTrainedModel.resize_position_embeddings(new_num_position_embeddings)

resize the model position embeddings if necessary

Source code in mindnlp/transformers/modeling_utils.py

def resize_position_embeddings(self, new_num_position_embeddings: int):
    """
    resize the model position embeddings if necessary
    """
    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__}"
    )

mindnlp.transformers.modeling_utils.PreTrainedModel.resize_token_embeddings(new_num_tokens=None, pad_to_multiple_of=None)

Resizes input token embeddings matrix of the model if new_num_tokens != config.vocab_size.

Takes care of tying weights embeddings afterwards if the model class has a tie_weights() method.

PARAMETER	DESCRIPTION
new_num_tokens	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 torch.nn.Embedding module of the model without doing anything. TYPE: `int`, *optional* DEFAULT: None

RETURNS	DESCRIPTION
Embedding	torch.nn.Embedding: Pointer to the input tokens Embeddings Module of the model.

Source code in mindnlp/transformers/modeling_utils.py

def resize_token_embeddings(
    self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None
) -> nn.Embedding:
    """
    Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.

    Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.

    Arguments:
        new_num_tokens (`int`, *optional*):
            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 `torch.nn.Embedding` module of the model without doing anything.

    Returns:
        `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
    """
    model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of)
    if new_num_tokens is None and pad_to_multiple_of is None:
        return model_embeds
    # Update base model and current model config
    self.config.vocab_size = model_embeds.weight.shape[0]
    self.vocab_size = model_embeds.weight.shape[0]
    # Tie weights again if needed
    self.tie_weights()

    return model_embeds

mindnlp.transformers.modeling_utils.PreTrainedModel.resize_tokenizer_embeddings(new_num_tokens)

Obtain a new embedding layer or use the original one without updating it.

Source code in mindnlp/transformers/modeling_utils.py

def resize_tokenizer_embeddings(self, new_num_tokens):
    """
    Obtain a new embedding layer or use the original one without updating it.
    """
    old_embeddings = self.get_input_embeddings()
    new_embeddings = self._get_resized_embeddings(
        old_embeddings, new_num_tokens)
    self.set_input_embeddings(new_embeddings)
    return self.get_input_embeddings()

mindnlp.transformers.modeling_utils.PreTrainedModel.save(save_dir)

Save a model and its configuration file to a directory, so that it can be re-loaded using the :func:PreTrainedModel.from_pretrained`` class method.

PARAMETER	DESCRIPTION
save_dir	directory to which to save.

Source code in mindnlp/transformers/modeling_utils.py

def save(self, save_dir):
    """
    Save a model and its configuration file to a directory, so that
    it can be re-loaded using the `:func:`PreTrainedModel.from_pretrained`` class method.

    Arguments:
        save_dir: directory to which to save.
    """
    if os.path.isfile(save_dir):
        logger.error(f"Provided path ({save_dir}) should be a directory, not a file")
        return
    os.makedirs(save_dir, exist_ok=True)

    # Only save the model itself if we are using distributed training
    model_to_save = self.cell if hasattr(self, "cell") else self

    # Attach architecture to the config
    model_to_save.config.architectures = [model_to_save.__class__.__name__]

    # If we save using the predefined names, we can load using `from_pretrained`
    output_model_file = os.path.join(save_dir, WEIGHTS_NAME)
    save_checkpoint(model_to_save, output_model_file)

    logger.info(f"Model weights saved in {output_model_file}")

mindnlp.transformers.modeling_utils.PreTrainedModel.save_pretrained(save_directory, is_main_process=True, state_dict=None, save_function=mindspore.save_checkpoint, max_shard_size='5GB', safe_serialization=True, variant=None, **kwargs)

Save a model and its configuration file to a directory, so that it can be re-loaded using the [~PreTrainedModel.from_pretrained] class method.

PARAMETER	DESCRIPTION
save_directory	Directory to which to save. Will be created if it doesn't exist. TYPE: `str` or `os.PathLike`
is_main_process	Whether the process calling this is the main process or not. Useful when in distributed training like TPUs and need to call this function on all processes. In this case, set is_main_process=True only on the main process to avoid race conditions. TYPE: `bool`, *optional*, defaults to `True` DEFAULT: True
state_dict	The state dictionary of the model to save. Will default to self.state_dict(), but can be used to only save parts of the model or if special precautions need to be taken when recovering the state dictionary of a model (like when using model parallelism). TYPE: nested dictionary of `torch.Tensor` DEFAULT: None
save_function	The function to use to save the state dictionary. Useful on distributed training like TPUs when one need to replace torch.save by another method. TYPE: `Callable` DEFAULT: save_checkpoint
max_shard_size	The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size lower than this size. If expressed as a string, needs to be digits followed by a unit (like "5MB"). We default it to 5GB in order for models to be able to run easily on free-tier google colab instances without CPU OOM issues. If a single weight of the model is bigger than max_shard_size, it will be in its own checkpoint shard which will be bigger than max_shard_size. TYPE: `int` or `str`, *optional*, defaults to `"5GB"` DEFAULT: '5GB'
variant	If specified, weights are saved in the format pytorch_model..bin. TYPE: `str`, *optional* DEFAULT: None
save_peft_format	For backward compatibility with PEFT library, in case adapter weights are attached to the model, all keys of the state dict of adapters needs to be pre-pended with base_model.model. Advanced users can disable this behaviours by setting save_peft_format to False. TYPE: `bool`, *optional*, defaults to `True`
kwargs	Additional key word arguments passed along to the [~utils.PushToHubMixin.push_to_hub] method. TYPE: `Dict[str, Any]`, *optional* DEFAULT: {}

Source code in mindnlp/transformers/modeling_utils.py

def save_pretrained(
    self,
    save_directory: Union[str, os.PathLike],
    is_main_process: bool = True,
    state_dict: Optional[dict] = None,
    save_function: Callable = mindspore.save_checkpoint,
    max_shard_size: Union[int, str] = "5GB",
    safe_serialization: bool = True,
    variant: Optional[str] = None,
    **kwargs,
):
    """
    Save a model and its configuration file to a directory, so that it can be re-loaded using the
    [`~PreTrainedModel.from_pretrained`] class method.

    Arguments:
        save_directory (`str` or `os.PathLike`):
            Directory to which to save. Will be created if it doesn't exist.
        is_main_process (`bool`, *optional*, defaults to `True`):
            Whether the process calling this is the main process or not. Useful when in distributed training like
            TPUs and need to call this function on all processes. In this case, set `is_main_process=True` only on
            the main process to avoid race conditions.
        state_dict (nested dictionary of `torch.Tensor`):
            The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only
            save parts of the model or if special precautions need to be taken when recovering the state dictionary
            of a model (like when using model parallelism).
        save_function (`Callable`):
            The function to use to save the state dictionary. Useful on distributed training like TPUs when one
            need to replace `torch.save` by another method.
        max_shard_size (`int` or `str`, *optional*, defaults to `"5GB"`):
            The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size
            lower than this size. If expressed as a string, needs to be digits followed by a unit (like `"5MB"`).
            We default it to 5GB in order for models to be able to run easily on free-tier google colab instances
            without CPU OOM issues.

            <Tip warning={true}>

            If a single weight of the model is bigger than `max_shard_size`, it will be in its own checkpoint shard
            which will be bigger than `max_shard_size`.

            </Tip>
        variant (`str`, *optional*):
            If specified, weights are saved in the format pytorch_model.<variant>.bin.
        save_peft_format (`bool`, *optional*, defaults to `True`):
            For backward compatibility with PEFT library, in case adapter weights are attached to the model, all
            keys of the state dict of adapters needs to be pre-pended with `base_model.model`. Advanced users can
            disable this behaviours by setting `save_peft_format` to `False`.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
    """
    if os.path.isfile(save_directory):
        logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
        return

    os.makedirs(save_directory, exist_ok=True)

    # Only save the model itself if we are using distributed training
    model_to_save = self

    # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
    # we currently don't use this setting automatically, but may start to use with v5
    dtype = get_parameter_dtype(model_to_save)
    model_to_save.config.ms_dtype = str(dtype).lower()

    # Attach architecture to the config
    model_to_save.config.architectures = [model_to_save.__class__.__name__]

    # Save the config
    if is_main_process:
        model_to_save.config.save_pretrained(save_directory)
        if self.can_generate():
            model_to_save.generation_config.save_pretrained(save_directory)

    # Save the model
    if state_dict is None:
        state_dict = model_to_save.parameters_dict()

    # Handle the case where some state_dict keys shouldn't be saved
    if self._keys_to_ignore_on_save is not None:
        for ignore_key in self._keys_to_ignore_on_save:
            if ignore_key in state_dict.keys():
                del state_dict[ignore_key]

    # Shard the model if it is too big.
    # weights_name = _add_variant(WEIGHTS_NAME, variant)
    weights_name = SAFE_WEIGHTS_NAME if safe_serialization else WEIGHTS_NAME

    shards, index = shard_checkpoint(state_dict, max_shard_size=max_shard_size, weights_name=weights_name)

    # Clean the folder from a previous save
    for filename in os.listdir(save_directory):
        full_filename = os.path.join(save_directory, filename)
        # If we have a shard file that is not going to be replaced, we delete it, but only from the main process
        # in distributed settings to avoid race conditions.
        weights_no_suffix = weights_name.replace(".bin", "").replace(".safetensors", "")

        # make sure that file to be deleted matches format of sharded file, e.g. pytorch_model-00001-of-00005
        filename_no_suffix = filename.replace(".bin", "").replace(".safetensors", "")
        reg = re.compile(r"(.*?)-\d{5}-of-\d{5}")

        if (
            filename.startswith(weights_no_suffix)
            and os.path.isfile(full_filename)
            and filename not in shards
            and is_main_process
            and reg.fullmatch(filename_no_suffix) is not None
        ):
            os.remove(full_filename)

    # Save the model
    for shard_file, shard in shards.items():
        if safe_serialization:
            # At some point we will need to deal better with save_function (used for TPU and other distributed
            # joyfulness), but for now this enough.
            safe_save_file(shard, os.path.join(save_directory, shard_file), metadata={"format": "np"})
        else:
            save_function(shard, os.path.join(save_directory, shard_file))

    if index is None:
        path_to_weights = os.path.join(save_directory, _add_variant(WEIGHTS_NAME, variant))
        logger.info(f"Model weights saved in {path_to_weights}")
    else:
        save_index_file = SAFE_WEIGHTS_INDEX_NAME if safe_serialization else WEIGHTS_INDEX_NAME
        save_index_file = os.path.join(save_directory, _add_variant(save_index_file, variant))
        # Save the index as well
        with open(save_index_file, "w", encoding="utf-8") as f:
            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
            f.write(content)
        logger.info(
            f"The model is bigger than the maximum size per checkpoint ({max_shard_size}) and is going to be "
            f"split in {len(shards)} checkpoint shards. You can find where each parameters has been saved in the "
            f"index located at {save_index_file}."
        )

mindnlp.transformers.modeling_utils.PreTrainedModel.set_input_embeddings(new_embeddings)

Set model's input embeddings.

PARAMETER	DESCRIPTION
value

Source code in mindnlp/transformers/modeling_utils.py

def set_input_embeddings(self, new_embeddings: nn.Module):
    """
    Set model's input embeddings.

    Args:
        value (:obj:`nn.Module`): A mindspore cell mapping vocabulary to hidden states.
    """
    base_model = getattr(self, self.base_model_prefix, self)
    if base_model is not self:
        return base_model.set_input_embeddings(new_embeddings)
    raise NotImplementedError

mindnlp.transformers.modeling_utils.PreTrainedModel.set_output_embeddings(new_embeddings)

Set model's output embeddings.

PARAMETER	DESCRIPTION
value

Source code in mindnlp/transformers/modeling_utils.py

def set_output_embeddings(self, new_embeddings: nn.Module):
    """
    Set model's output embeddings.

    Args:
        value (:obj:`nn.Module`): A mindspore cell mapping vocabulary to hidden states.
    """
    base_model = getattr(self, self.base_model_prefix, self)
    if base_model is not self:
        return base_model.set_output_embeddings(new_embeddings)
    raise NotImplementedError

mindnlp.transformers.modeling_utils.PreTrainedModel.tie_weights()

Make sure we are sharing the input and output embeddings. If you need this feature, you need to get it yourself output Add the output you need to add to the embeddings function_ Embedding layer, otherwise you cannot

Source code in mindnlp/transformers/modeling_utils.py

def tie_weights(self):
    """
    Make sure we are sharing the input and output embeddings.
    If you need this feature,
    you need to get it yourself output Add the output you need to add to the embeddings function_ Embedding layer,
    otherwise you cannot
    """
    if getattr(self.config, "tie_word_embeddings", True):
        output_embeddings = self.get_output_embeddings() # pylint: disable=assignment-from-none
        if output_embeddings is not None:
            self._tie_or_clone_weights(
                output_embeddings, self.get_input_embeddings())

    if getattr(self.config, "is_encoder_decoder", False) and getattr(self.config, "tie_encoder_decoder", False):
        if hasattr(self, self.base_model_prefix):
            self = getattr(self, self.base_model_prefix) # pylint: disable=self-cls-assignment
        self._tie_encoder_decoder_weights(
            self.encoder, self.decoder, self.base_model_prefix)

    for _, cell in self.cells_and_names():
        if hasattr(cell, "_tie_weights"):
            cell._tie_weights()

mindnlp.transformers.modeling_utils.PreTrainedModel.trainable_params(recurse=True)

fix duplicated weights

Source code in mindnlp/transformers/modeling_utils.py

def trainable_params(self, recurse=True):
    """
    fix duplicated weights
    """
    return list(set(filter(lambda x: x.requires_grad, self.get_parameters(expand=recurse))))

mindnlp.transformers.modeling_utils.PreTrainedModel.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)

Shows a one-time warning if the input_ids appear to contain padding and no attention mask was given.

Source code in mindnlp/transformers/modeling_utils.py

def warn_if_padding_and_no_attention_mask(self, input_ids, attention_mask):
    """
    Shows a one-time warning if the input_ids appear to contain padding and no attention mask was given.
    """
    if (attention_mask is not None) or (self.config.pad_token_id is None):
        return

    # Check only the first and last input IDs to reduce overhead.
    # if self.config.pad_token_id in input_ids[:, [-1, 0]]:
    if ops.contains(input_ids[:, [-1, 0]], self.config.pad_token_id):
        warn_string = (
            "We strongly recommend passing in an `attention_mask` since your input_ids may be padded."
        )

        # If the pad token is equal to either BOS, EOS, or SEP, we do not know whether the user should use an
        # attention_mask or not. In this case, we should still show a warning because this is a rare case.
        if (
            (self.config.bos_token_id is not None and self.config.bos_token_id == self.config.pad_token_id)
            or (self.config.eos_token_id is not None and self.config.eos_token_id == self.config.pad_token_id)
            or (self.config.sep_token_id is not None and self.config.sep_token_id == self.config.pad_token_id)
        ):
            warn_string += (
                f"\nYou may ignore this warning if your `pad_token_id` ({self.config.pad_token_id}) is identical "
                f"to the `bos_token_id` ({self.config.bos_token_id}), `eos_token_id` ({self.config.eos_token_id}), "
                f"or the `sep_token_id` ({self.config.sep_token_id}), and your input is not padded."
            )

        logger.warning(warn_string)