/paddlenlp/transformers/utils.py

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

import contextlib
import functools
import hashlib
import importlib
import inspect
import os
import re
import shutil
import sys
import warnings
from contextlib import ExitStack
from io import StringIO
from pathlib import Path
from typing import TYPE_CHECKING, ContextManager, List, Optional, Type, Union

from filelock import FileLock

from paddlenlp import __version__
from paddlenlp.utils.downloader import (
    COMMUNITY_MODEL_PREFIX,
    download_check,
    get_path_from_url_with_filelock,
    is_url,
    url_file_exists,
)

if TYPE_CHECKING:
    from paddlenlp.transformers import PretrainedModel

import numpy as np
import paddle
import tqdm
from huggingface_hub import hf_hub_download, try_to_load_from_cache
from huggingface_hub.utils import EntryNotFoundError
from paddle.common_ops_import import convert_dtype
from paddle.nn import Layer
from requests.exceptions import HTTPError

from paddlenlp.utils.env import HF_CACHE_HOME, MODEL_HOME
from paddlenlp.utils.import_utils import import_module
from paddlenlp.utils.log import logger

from ..utils.download import resolve_file_path

# TODO(@zewu): upgrade aistudio to the newest version
try:
    from .aistudio_utils import aistudio_download
except:
    aistudio_download = None

HUGGINGFACE_CO_RESOLVE_ENDPOINT = "https://huggingface.co"


def convert_ndarray_dtype(np_array: np.ndarray, target_dtype: str) -> np.ndarray:
    """convert ndarray

    Args:
        np_array (np.ndarray): numpy ndarray instance
        target_dtype (str): the target dtype

    Returns:
        np.ndarray: converted numpy ndarray instance
    """
    source_dtype = convert_dtype(np_array.dtype)
    if source_dtype == "uint16" or target_dtype == "bfloat16":
        tensor = paddle.to_tensor(np_array)
        tensor = paddle.cast(tensor, target_dtype)
        return tensor.cpu().numpy()

        # TODO(wj-Mcat): device_guard will slow the converting
        # with device_guard("cpu"):
        #     tensor = paddle.to_tensor(np_array)
        #     tensor = paddle.cast(tensor, target_dtype)
        # return tensor.cpu().numpy()

    if target_dtype == "bfloat16":
        target_dtype = "uint16"

    return np_array.astype(target_dtype)


def convert_to_dict_message(conversation: List[List[str]]):
    """Convert the list of chat messages to a role dictionary chat messages."""
    conversations = []
    for index, item in enumerate(conversation):
        assert 1 <= len(item) <= 2, "Each Rounds in conversation should have 1 or 2 elements."
        if isinstance(item[0], str):
            conversations.append({"role": "user", "content": item[0]})
            if len(item) == 2 and isinstance(item[1], str):
                conversations.append({"role": "assistant", "content": item[1]})
            else:
                # If there is only one element in item, it must be the last round.
                # If it is not the last round, it must be an error.
                if index != len(conversation) - 1:
                    raise ValueError(f"Round {index} has error round")
        else:
            raise ValueError("Each round in list should be string")
    return conversations


def get_scale_by_dtype(dtype: str = None, return_positive: bool = True) -> float:
    """get scale value by dtype

    Args:
        dtype (str): the string dtype value

    Returns:
        float: the scale value
    """
    if dtype is None:
        dtype = paddle.get_default_dtype()

    dtype = convert_dtype(dtype)
    scale_value = 1e6

    # TODO(wj-Mcaf): support int8, int4 dtypes later
    if dtype == "float16":
        scale_value = 1e4

    if return_positive:
        return scale_value
    return -1 * scale_value


def fn_args_to_dict(func, *args, **kwargs):
    """
    Inspect function `func` and its arguments for running, and extract a
    dict mapping between argument names and keys.
    """
    if hasattr(inspect, "getfullargspec"):
        (spec_args, spec_varargs, spec_varkw, spec_defaults, _, _, _) = inspect.getfullargspec(func)
    else:
        (spec_args, spec_varargs, spec_varkw, spec_defaults) = inspect.getargspec(func)
    # add positional argument values
    init_dict = dict(zip(spec_args, args))
    # add default argument values
    kwargs_dict = dict(zip(spec_args[-len(spec_defaults) :], spec_defaults)) if spec_defaults else {}
    for k in list(kwargs_dict.keys()):
        if k in init_dict:
            kwargs_dict.pop(k)
    kwargs_dict.update(kwargs)
    init_dict.update(kwargs_dict)
    return init_dict


def adapt_stale_fwd_patch(self, name, value):
    """
    Since there are some monkey patches for forward of PretrainedModel, such as
    model compression, we make these patches compatible with the latest forward
    method.
    """
    if name == "forward":
        # NOTE(guosheng): In dygraph to static, `layer.forward` would be patched
        # by an instance of `StaticFunction`. And use string compare to avoid to
        # import fluid.
        if type(value).__name__.endswith("StaticFunction") or self.forward.__class__.__name__.endswith(
            "StaticFunction"
        ):
            return value
        if type(value).__name__.endswith("WeakMethod") or self.forward.__class__.__name__.endswith("WeakMethod"):
            return value

        # NOTE(changwenbin & zhoukangkang):
        # When use model = paddle.incubate.jit.inference(model), it reportes errors, we fix it here.
        # is_inference_mode API is only available in PaddlePaddle develop，so we add a try except.
        try:
            from paddle.incubate.jit import is_inference_mode

            if is_inference_mode(value):
                return value
        except:
            pass

        if hasattr(inspect, "getfullargspec"):
            (
                patch_spec_args,
                patch_spec_varargs,
                patch_spec_varkw,
                patch_spec_defaults,
                _,
                _,
                _,
            ) = inspect.getfullargspec(value)
            (spec_args, spec_varargs, spec_varkw, spec_defaults, _, _, _) = inspect.getfullargspec(self.forward)
        else:
            (patch_spec_args, patch_spec_varargs, patch_spec_varkw, patch_spec_defaults) = inspect.getargspec(value)
            (spec_args, spec_varargs, spec_varkw, spec_defaults) = inspect.getargspec(self.forward)
        new_args = [
            arg
            for arg in ("output_hidden_states", "output_attentions", "return_dict")
            if arg not in patch_spec_args and arg in spec_args
        ]

        if new_args:
            if self.__module__.startswith("paddlenlp"):
                warnings.warn(
                    f"The `forward` method of {self.__class__ if isinstance(self, Layer) else self} is patched and the patch "
                    "might be based on an old version which missing some "
                    f"arguments compared with the latest, such as {new_args}. "
                    "We automatically add compatibility on the patch for "
                    "these arguments, and maybe the patch should be updated."
                )
            else:
                warnings.warn(
                    f"The `forward` method of {self.__class__ if isinstance(self, Layer) else self} "
                    "is patched and the patch might be conflict with patches made "
                    f"by paddlenlp which seems have more arguments such as {new_args}. "
                    "We automatically add compatibility on the patch for "
                    "these arguments, and maybe the patch should be updated."
                )
            if isinstance(self, Layer) and inspect.isfunction(value):

                @functools.wraps(value)
                def wrap_fwd(*args, **kwargs):
                    for arg in new_args:
                        kwargs.pop(arg, None)
                    return value(self, *args, **kwargs)

            else:

                @functools.wraps(value)
                def wrap_fwd(*args, **kwargs):
                    for arg in new_args:
                        kwargs.pop(arg, None)
                    return value(*args, **kwargs)

            return wrap_fwd
    return value


class InitTrackerMeta(type(Layer)):
    """
    This metaclass wraps the `__init__` method of a class to add `init_config`
    attribute for instances of that class, and `init_config` use a dict to track
    the initial configuration. If the class has `_pre_init` or `_post_init`
    method, it would be hooked before or after `__init__` and called as
    `_pre_init(self, init_fn, init_args)` or `_post_init(self, init_fn, init_args)`.
    Since InitTrackerMeta would be used as metaclass for pretrained model classes,
    which always are Layer and `type(Layer)` is not `type`, thus use `type(Layer)`
    rather than `type` as base class for it to avoid inheritance metaclass
    conflicts.
    """

    def __init__(cls, name, bases, attrs):
        init_func = cls.__init__
        # If attrs has `__init__`, wrap it using accessible `_pre_init, _post_init`.
        # Otherwise, no need to wrap again since the super cls has been wrapped.
        # TODO: remove reduplicated tracker if using super cls `__init__`
        pre_init_func = getattr(cls, "_pre_init", None) if "__init__" in attrs else None
        post_init_func = getattr(cls, "_post_init", None) if "__init__" in attrs else None
        cls.__init__ = InitTrackerMeta.init_and_track_conf(init_func, pre_init_func, post_init_func)
        super(InitTrackerMeta, cls).__init__(name, bases, attrs)

    @staticmethod
    def init_and_track_conf(init_func, pre_init_func=None, post_init_func=None):
        """
        wraps `init_func` which is `__init__` method of a class to add `init_config`
        attribute for instances of that class.
        Args:
            init_func (callable): It should be the `__init__` method of a class.
                warning: `self` always is the class type of down-stream model, eg: BertForTokenClassification
            pre_init_func (callable, optional): If provided, it would be hooked after
                `init_func` and called as `pre_init_func(self, init_func, *init_args, **init_args)`.
                Default None.
            post_init_func (callable, optional): If provided, it would be hooked after
                `init_func` and called as `post_init_func(self, init_func, *init_args, **init_args)`.
                Default None.

        Returns:
            function: the wrapped function
        """

        @functools.wraps(init_func)
        def __impl__(self, *args, **kwargs):
            # registered helper by `pre_init_func`
            if pre_init_func:
                pre_init_func(self, init_func, *args, **kwargs)
            # keep full configuration
            init_func(self, *args, **kwargs)
            # registered helper by `post_init_func`
            if post_init_func:
                post_init_func(self, init_func, *args, **kwargs)
            self.init_config = kwargs
            if args:
                kwargs["init_args"] = args
            kwargs["init_class"] = self.__class__.__name__

        return __impl__

    def __setattr__(self, name, value):
        value = adapt_stale_fwd_patch(self, name, value)
        return super(InitTrackerMeta, self).__setattr__(name, value)


def param_in_func(func, param_field: str) -> bool:
    """check if the param_field is in `func` method, eg: if the `bert` param is in `__init__` method

    Args:
        cls (type): the class of PretrainedModel
        param_field (str): the name of field

    Returns:
        bool: the result of existence
    """

    if hasattr(inspect, "getfullargspec"):
        result = inspect.getfullargspec(func)
    else:
        result = inspect.getargspec(func)

    return param_field in result[0]


def resolve_cache_dir(from_hf_hub: bool, from_aistudio: bool, cache_dir: Optional[str] = None) -> str:
    """resolve cache dir for PretrainedModel and PretrainedConfig

    Args:
        from_hf_hub (bool): if load from huggingface hub
        cache_dir (str): cache_dir for models
    """
    if cache_dir is not None:
        return cache_dir
    if from_aistudio:
        return None
    if from_hf_hub:
        return HF_CACHE_HOME
    return MODEL_HOME


def find_transformer_model_type(model_class: Type) -> str:
    """get the model type from module name,
        eg:
            BertModel -> bert,
            RobertaForTokenClassification -> roberta

    Args:
        model_class (Type): the class of model

    Returns:
        str: the type string
    """
    from paddlenlp.transformers import PretrainedModel

    default_model_type = ""

    if not issubclass(model_class, PretrainedModel):
        return default_model_type

    module_name: str = model_class.__module__
    if not module_name.startswith("paddlenlp.transformers."):
        return default_model_type

    tokens = module_name.split(".")
    if len(tokens) < 3:
        return default_model_type

    return tokens[2]


def find_transformer_model_class_by_name(model_name: str) -> Optional[Type[PretrainedModel]]:
    """find transformer model_class by name

    Args:
        model_name (str): the string of class name

    Returns:
        Optional[Type[PretrainedModel]]: optional pretrained-model class
    """
    transformer_module = import_module("paddlenlp.transformers")

    for obj_name in dir(transformer_module):
        if obj_name.startswith("_"):
            continue
        obj = getattr(transformer_module, obj_name, None)
        if obj is None:
            continue

        name = getattr(obj, "__name__", None)
        if name is None:
            continue

        if name == model_name:
            return obj
    logger.debug(f"can not find model_class<{model_name}>")
    return None


def convert_file_size_to_int(size: Union[int, str]):
    """
    Converts a size expressed as a string with digits an unit (like `"5MB"`) to an integer (in bytes).
    Args:
        size (`int` or `str`): The size to convert. Will be directly returned if an `int`.
    Example:
    ```py
    >>> convert_file_size_to_int("1MiB")
    1048576
    ```
    """
    if isinstance(size, int):
        return size
    if size.upper().endswith("GIB"):
        return int(size[:-3]) * (2**30)
    if size.upper().endswith("MIB"):
        return int(size[:-3]) * (2**20)
    if size.upper().endswith("KIB"):
        return int(size[:-3]) * (2**10)
    if size.upper().endswith("GB"):
        int_size = int(size[:-2]) * (10**9)
        return int_size // 8 if size.endswith("b") else int_size
    if size.upper().endswith("MB"):
        int_size = int(size[:-2]) * (10**6)
        return int_size // 8 if size.endswith("b") else int_size
    if size.upper().endswith("KB"):
        int_size = int(size[:-2]) * (10**3)
        return int_size // 8 if size.endswith("b") else int_size
    raise ValueError("`size` is not in a valid format. Use an integer followed by the unit, e.g., '5GB'.")


def paddlenlp_hub_download(
    repo_id: str,
    filename: str,
    *,
    subfolder: Optional[str] = None,
    cache_dir: Union[str, Path, None] = None,
    pretrained_model_name_or_path: str = None,
) -> str:
    if subfolder is None:
        subfolder = ""
    if pretrained_model_name_or_path is not None and is_url(repo_id):
        cache_dir = os.path.join(cache_dir, pretrained_model_name_or_path, subfolder)
    else:
        cache_dir = os.path.join(cache_dir, repo_id, subfolder)

    # check in cache_dir
    weight_file_path = os.path.join(cache_dir, filename)

    if os.path.exists(weight_file_path):
        logger.info(f"Already cached {weight_file_path}")
        return weight_file_path

    # Download from custom model url
    if is_url(repo_id):
        # check whether the target file exist in the community bos server
        if url_file_exists(repo_id):
            logger.info(f"Downloading {repo_id}")
            weight_file_path = get_path_from_url_with_filelock(repo_id, cache_dir)
            # # check the downloaded weight file and registered weight file name
            download_check(repo_id, "paddlenlp_hub_download")

            # make sure that model states names: model_states.pdparams
            new_weight_file_path = os.path.join(os.path.split(weight_file_path)[0], filename)

            if weight_file_path != new_weight_file_path:
                # create lock file, which is empty, under the `LOCK_FILE_HOME` directory.
                lock_file_name = hashlib.md5((repo_id + cache_dir).encode("utf-8")).hexdigest()
                # create `.lock` private directory in the cache dir
                lock_file_path = os.path.join(cache_dir, ".lock", lock_file_name)

                with FileLock(lock_file_path):
                    if not os.path.exists(new_weight_file_path):
                        shutil.move(weight_file_path, new_weight_file_path)

                weight_file_path = new_weight_file_path

            return weight_file_path

        return None

    # find in community repo
    url_list = [COMMUNITY_MODEL_PREFIX, repo_id, filename]
    if subfolder != "":
        url_list.insert(2, subfolder)
    community_model_file_path = "/".join(url_list)
    assert is_url(community_model_file_path)

    # check whether the target file exist in the community bos server
    if url_file_exists(community_model_file_path):
        logger.info(f"Downloading {community_model_file_path}")
        weight_file_path = get_path_from_url_with_filelock(community_model_file_path, cache_dir)
        # # check the downloaded weight file and registered weight file name
        download_check(community_model_file_path, "paddlenlp_hub_download")
        return weight_file_path

    return None


# Return value when trying to load a file from cache but the file does not exist in the distant repo.
_CACHED_NO_EXIST = object()


def cached_file(
    path_or_repo_id: Union[str, os.PathLike],
    filename: str,
    cache_dir: Optional[Union[str, os.PathLike]] = None,
    subfolder: str = "",
    from_aistudio: bool = False,
    _raise_exceptions_for_missing_entries: bool = True,
    _raise_exceptions_for_connection_errors: bool = True,
    pretrained_model_name_or_path=None,
) -> str:
    """
    Tries to locate a file in a local folder and repo, downloads and cache it if necessary.
    Args:
        path_or_repo_id (`str` or `os.PathLike`):
            This can be either:
            - a string, the *model id* of a model repo on huggingface.co.
            - a path to a *directory* potentially containing the file.
        filename (`str`):
            The name of the file to locate in `path_or_repo`.
        cache_dir (`str` or `os.PathLike`, *optional*):
            Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
            cache should not be used.
        subfolder (`str`, *optional*, defaults to `""`):
            In case the relevant files are located inside a subfolder of the model repo on huggingface.co, you can
            specify the folder name here.

    Returns:
        `Optional[str]`: Returns the resolved file (to the cache folder if downloaded from a repo).
    Examples:
    ```python
    # Download a model weight from the Hub and cache it.
    model_weights_file = cached_file("bert-base-uncased", "pytorch_model.bin")
    ```
    """

    if subfolder is None:
        subfolder = ""

    path_or_repo_id = str(path_or_repo_id)
    full_filename = os.path.join(subfolder, filename)
    if os.path.isdir(path_or_repo_id):
        resolved_file = os.path.join(os.path.join(path_or_repo_id, subfolder), filename)
        if not os.path.isfile(resolved_file):
            if _raise_exceptions_for_missing_entries:
                raise EnvironmentError(
                    f"{path_or_repo_id} does not appear to have a file named {full_filename}. Checkout "
                    f"'https://huggingface.co/{path_or_repo_id}/' for available files."
                )
            else:
                return None
        return resolved_file

    if cache_dir is not None and isinstance(cache_dir, Path):
        cache_dir = str(cache_dir)

    if from_aistudio:
        try:
            resolved_file = aistudio_download(
                repo_id=path_or_repo_id, filename=filename, subfolder=subfolder, cache_dir=cache_dir
            )
        except:
            resolved_file = None
    else:
        # if cache_dir is None:
        #     cache_dir = os.path.join(MODEL_HOME, ".cache")
        try:
            # Load from URL or cache if already cached
            resolved_file = paddlenlp_hub_download(
                path_or_repo_id,
                filename,
                subfolder=None if len(subfolder) == 0 else subfolder,
                # revision=revision,
                cache_dir=cache_dir,
                pretrained_model_name_or_path=pretrained_model_name_or_path,
            )
        except HTTPError as err:
            # First we try to see if we have a cached version (not up to date):
            resolved_file = try_to_load_from_cache(path_or_repo_id, full_filename, cache_dir=cache_dir)
            if resolved_file is not None and resolved_file != _CACHED_NO_EXIST:
                return resolved_file
            if not _raise_exceptions_for_connection_errors:
                return None

            raise EnvironmentError(
                f"There was a specific connection error when trying to load {path_or_repo_id}:\n{err}"
            )

    return resolved_file


def cached_file_for_hf_hub(
    path_or_repo_id: Union[str, os.PathLike],
    filename: str,
    cache_dir: Optional[Union[str, os.PathLike]] = None,
    subfolder: str = "",
    _raise_exceptions_for_missing_entries: bool = True,
):

    if subfolder is None:
        subfolder = ""

    path_or_repo_id = str(path_or_repo_id)
    full_filename = os.path.join(subfolder, filename)
    if os.path.isdir(path_or_repo_id):
        resolved_file = os.path.join(os.path.join(path_or_repo_id, subfolder), filename)
        if not os.path.isfile(resolved_file):
            if _raise_exceptions_for_missing_entries:
                raise EnvironmentError(
                    f"{path_or_repo_id} does not appear to have a file named {full_filename}. Checkout "
                    f"'https://huggingface.co/{path_or_repo_id}' for available files."
                )
            else:
                return None
        return resolved_file

    if cache_dir is None:
        cache_dir = os.path.join(MODEL_HOME, ".cache")
    if isinstance(cache_dir, Path):
        cache_dir = str(cache_dir)

    try:
        # Load from URL or cache if already cached
        download_check(path_or_repo_id, full_filename, addition="from_hf_hub")
        resolved_file = hf_hub_download(
            repo_id=path_or_repo_id,
            filename=filename,
            cache_dir=cache_dir,
            subfolder=subfolder,
            library_name="PaddleNLP",
            library_version=__version__,
        )
        return resolved_file
    except Exception as e:
        print(e)
        msg = f"""
            {path_or_repo_id} is not a local folder and is not a valid model identifier "
            "listed on 'https://huggingface.co/models' If this is a private repository, make sure to "
            "pass a token having permission to this repo with `use_auth_token` or log in with "
            "`huggingface-cli login` and pass `use_auth_token=True`.
            """
        if _raise_exceptions_for_missing_entries:
            raise EnvironmentError(msg)
        else:
            logger.info(msg)
            return None


def get_checkpoint_shard_files(
    pretrained_model_name_or_path,
    index_filename,
    cache_dir=None,
    subfolder="",
    from_aistudio=False,
    from_hf_hub=False,
):
    """
    For a given model:
    - download and cache all the shards of a sharded checkpoint if `pretrained_model_name_or_path` is a model ID on the
      Hub
    - returns the list of paths to all the shards, as well as some metadata.
    For the description of each arg, see [`PretrainedModel.from_pretrained`]. `index_filename` is the full path to the
    index (downloaded and cached if `pretrained_model_name_or_path` is a model ID on the Hub).
    """

    import json

    if not os.path.isfile(index_filename):
        raise ValueError(f"Can't find a checkpoint index ({index_filename}) in {pretrained_model_name_or_path}.")

    with open(index_filename, "r") as f:
        index = json.loads(f.read())

    shard_filenames = sorted(set(index["weight_map"].values()))
    sharded_metadata = index["metadata"]
    sharded_metadata["all_checkpoint_keys"] = list(index["weight_map"].keys())
    sharded_metadata["weight_map"] = index["weight_map"].copy()

    file_map = {file: set() for file in shard_filenames}
    for weight, file in index["weight_map"].items():
        file_map[file].add(weight)

    sharded_metadata["file_map"] = file_map

    # First, let's deal with local folder.
    if os.path.isdir(pretrained_model_name_or_path):
        shard_filenames = [os.path.join(pretrained_model_name_or_path, subfolder, f) for f in shard_filenames]
        return shard_filenames, sharded_metadata

    # At this stage pretrained_model_name_or_path is a model identifier on the Hub
    cached_filenames = []
    # Check if the model is already cached or not. We only try the last checkpoint, this should cover most cases of
    # downloaded (if interrupted).
    last_shard = try_to_load_from_cache(
        pretrained_model_name_or_path,
        shard_filenames[-1],
        cache_dir=cache_dir,
    )

    show_progress_bar = last_shard is None
    for shard_filename in tqdm.tqdm(shard_filenames, desc="Downloading shards", disable=not show_progress_bar):
        try:
            cached_filename = resolve_file_path(
                pretrained_model_name_or_path,
                [shard_filename],
                subfolder,
                cache_dir=cache_dir,
                from_aistudio=from_aistudio,
                from_hf_hub=from_hf_hub,
            )
            assert (
                cached_filename is not None
            ), f"please make sure {shard_filename} under {pretrained_model_name_or_path}"
        # We have already dealt with RepositoryNotFoundError and RevisionNotFoundError when getting the index, so
        # we don't have to catch them here.
        except EntryNotFoundError:
            raise EnvironmentError(
                f"{pretrained_model_name_or_path} does not appear to have a file named {shard_filename} which is "
                "required according to the checkpoint index."
            )
        except HTTPError:
            raise EnvironmentError(
                f"We couldn't connect to '{HUGGINGFACE_CO_RESOLVE_ENDPOINT}' to load {shard_filename}. You should try"
                " again after checking your internet connection."
            )

        cached_filenames.append(cached_filename)

    return cached_filenames, sharded_metadata


def is_safetensors_available():
    return importlib.util.find_spec("safetensors") is not None


@contextlib.contextmanager
def device_guard(device="cpu", dev_id=0):
    origin_device = paddle.device.get_device()
    if device == "cpu":
        paddle.set_device(device)
    elif device in ["gpu", "xpu", "npu"]:
        paddle.set_device("{}:{}".format(device, dev_id))
    try:
        yield
    finally:
        paddle.set_device(origin_device)


def paddlenlp_load(path, map_location="cpu"):
    assert map_location in ["cpu", "gpu", "xpu", "npu", "numpy", "np"]
    if map_location in ["numpy", "np"]:
        return paddle.load(path, return_numpy=True)
    else:
        with device_guard(map_location):
            return paddle.load(path)
            # TODO(zhonghui03): the following code has problems when hot start optimizer checkpoint.
            if map_location == "cpu":
                from paddle.framework.io import (
                    _parse_every_object,
                    _to_LodTensor,
                    _transformed_from_lodtensor,
                )

                def _ndarray_to_tensor(obj, return_numpy=False):
                    if return_numpy:
                        return obj
                    if paddle.in_dynamic_mode():
                        return paddle.Tensor(obj, zero_copy=True)
                    else:
                        return _to_LodTensor(obj)

                state_dict = paddle.load(path, return_numpy=True)
                # Hack for zero copy for saving loading time. for paddle.load there need copy to create paddle.Tensor
                return _parse_every_object(state_dict, _transformed_from_lodtensor, _ndarray_to_tensor)

            else:
                return paddle.load(path)


def is_paddle_support_lazy_init():
    return hasattr(paddle, "LazyGuard")


class ContextManagers:
    """
    Wrapper for `contextlib.ExitStack` which enters a collection of context managers. Adaptation of `ContextManagers`
    in the `fastcore` library.
    """

    def __init__(self, context_managers: List[ContextManager]):
        self.context_managers = context_managers
        self.stack = ExitStack()

    def __enter__(self):
        for context_manager in self.context_managers:
            self.stack.enter_context(context_manager)

    def __exit__(self, *args, **kwargs):
        self.stack.__exit__(*args, **kwargs)


def use_hybrid_parallel():
    try:
        from paddle.distributed import fleet

        hcg = fleet.get_hybrid_communicate_group()
        return hcg
    except:
        return None


def optimizer_name_suffix():
    hcg = use_hybrid_parallel()
    if hcg is not None:
        name = []
        if hcg.get_model_parallel_world_size() > 1:
            name.append(f"tp{hcg.get_model_parallel_rank():0>2d}")
        if hcg.get_pipe_parallel_world_size() > 1:
            name.append(f"pp{hcg.get_stage_id():0>2d}")
        if hcg.get_sharding_parallel_world_size() > 1:
            name.append(f"shard{hcg.get_sharding_parallel_rank():0>2d}")

        return "_".join(name)
    else:
        return None


def weight_name_suffix():
    hcg = use_hybrid_parallel()
    if hcg is not None:
        name = []
        if hcg.get_model_parallel_world_size() > 1:
            name.append(f"tp{hcg.get_model_parallel_rank():0>2d}")
        if hcg.get_pipe_parallel_world_size() > 1:
            name.append(f"pp{hcg.get_stage_id():0>2d}")
        return "_".join(name)
    else:
        return None


def dtype_byte_size(dtype):
    """
    Returns the size (in bytes) occupied by one parameter of type `dtype`.

    Example:

    ```py
    >>> dtype_byte_size(paddle.float32)
    4
    ```
    """
    if dtype == paddle.bool:
        return 1 / 8
    if dtype == paddle.float8_e4m3fn or dtype == paddle.float8_e5m2:
        return 1
    bit_search = re.search(r"[^\d](\d+)$", str(dtype))
    if bit_search is None:
        raise ValueError(f"`dtype` is not a valid dtype: {dtype}.")
    bit_size = int(bit_search.groups()[0])
    return bit_size // 8


def apply_print_resets(buf):
    return re.sub(r"^.*\r", "", buf, 0, re.M)


class CaptureStd:
    """
    Context manager to capture:

        - stdout: replay it, clean it up and make it available via `obj.out`
        - stderr: replay it and make it available via `obj.err`

    Args:
        out (`bool`, *optional*, defaults to `True`): Whether to capture stdout or not.
        err (`bool`, *optional*, defaults to `True`): Whether to capture stderr or not.
        replay (`bool`, *optional*, defaults to `True`): Whether to replay or not.
            By default each captured stream gets replayed back on context's exit, so that one can see what the test was
            doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to
            disable this feature.

    Examples:

    ```python
    # to capture stdout only with auto-replay
    with CaptureStdout() as cs:
        print("Secret message")
    assert "message" in cs.out

    # to capture stderr only with auto-replay
    import sys

    with CaptureStderr() as cs:
        print("Warning: ", file=sys.stderr)
    assert "Warning" in cs.err

    # to capture both streams with auto-replay
    with CaptureStd() as cs:
        print("Secret message")
        print("Warning: ", file=sys.stderr)
    assert "message" in cs.out
    assert "Warning" in cs.err

    # to capture just one of the streams, and not the other, with auto-replay
    with CaptureStd(err=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    # but best use the stream-specific subclasses

    # to capture without auto-replay
    with CaptureStd(replay=False) as cs:
        print("Secret message")
    assert "message" in cs.out
    ```"""

    def __init__(self, out=True, err=True, replay=True):
        self.replay = replay

        if out:
            self.out_buf = StringIO()
            self.out = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.out_buf = None
            self.out = "not capturing stdout"

        if err:
            self.err_buf = StringIO()
            self.err = "error: CaptureStd context is unfinished yet, called too early"
        else:
            self.err_buf = None
            self.err = "not capturing stderr"

    def __enter__(self):
        if self.out_buf:
            self.out_old = sys.stdout
            sys.stdout = self.out_buf

        if self.err_buf:
            self.err_old = sys.stderr
            sys.stderr = self.err_buf

        return self

    def __exit__(self, *exc):
        if self.out_buf:
            sys.stdout = self.out_old
            captured = self.out_buf.getvalue()
            if self.replay:
                sys.stdout.write(captured)
            self.out = apply_print_resets(captured)

        if self.err_buf:
            sys.stderr = self.err_old
            captured = self.err_buf.getvalue()
            if self.replay:
                sys.stderr.write(captured)
            self.err = captured

    def __repr__(self):
        msg = ""
        if self.out_buf:
            msg += f"stdout: {self.out}\n"
        if self.err_buf:
            msg += f"stderr: {self.err}\n"
        return msg


def caculate_llm_per_token_flops(
    hidden_size,
    intermediate_size,
    layer_num,
    vocab_size,
    seq_length=None,
    recompute=False,
    recompute_granularity=None,
):

    # TFLOPs formula (from Equation 3 in Section 5.1 of https://arxiv.org/pdf/2104.04473.pdf).
    flops_per_transformer = 0
    flops_recompute_transformer = 0

    # qkvo matmul
    flops_qkvo_matmul = seq_length * hidden_size**2 * 4

    # [b,s,h] [b,h,s] bs^2h
    # [b,s,s] [b,s,h] bs^2h
    # q_states * k_states + attn_weight * v_states
    flops_core_attn = seq_length**2 * hidden_size * 2

    # swiglu, matmul + dot
    flops_ffn = seq_length * hidden_size * intermediate_size * 3 + seq_length * intermediate_size

    flops_per_transformer = flops_qkvo_matmul + flops_core_attn + flops_ffn
    if recompute:
        if recompute_granularity == "full":
            flops_recompute_transformer = flops_per_transformer
        if recompute_granularity == "full_attn":
            flops_recompute_transformer = flops_qkvo_matmul + flops_core_attn
        if recompute_granularity == "core_attn":
            flops_recompute_transformer = flops_core_attn

    # final loggits
    flops_loggits = seq_length * hidden_size * vocab_size

    # 2 for mul + add in matmul
    # 1 for forward, 2 for backwards since we caluate gradients for input_x and input_y
    return 2 * (layer_num * (flops_per_transformer * 3 + flops_recompute_transformer) + 3 * flops_loggits) / seq_length