huggingface-hub 0.20.2__py3-none-any.whl → 0.21.0__py3-none-any.whl

huggingface_hub/serialization/_tensorflow.py

@@ -0,0 +1,93 @@
+# Copyright 2024 The HuggingFace Team. 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.
+"""Contains tensorflow-specific helpers."""
+import math
+import re
+from typing import TYPE_CHECKING, Dict
+
+from ._base import MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
+
+
+if TYPE_CHECKING:
+    import tensorflow as tf
+
+
+def split_tf_state_dict_into_shards(
+    state_dict: Dict[str, "tf.Tensor"],
+    *,
+    filename_pattern: str = "tf_model{suffix}.h5",
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, Tensor]`):
+            The state dictionary to save.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"tf_model{suffix}.h5"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+    """
+    return split_state_dict_into_shards_factory(
+        state_dict,
+        max_shard_size=max_shard_size,
+        filename_pattern=filename_pattern,
+        get_tensor_size=get_tensor_size,
+    )
+
+
+def get_tensor_size(tensor: "tf.Tensor") -> int:
+    # Return `math.ceil` since dtype byte size can be a float (e.g., 0.125 for tf.bool).
+    # Better to overestimate than underestimate.
+    return math.ceil(tensor.numpy().size * _dtype_byte_size_tf(tensor.dtype))
+
+
+def _dtype_byte_size_tf(dtype) -> float:
+    """
+    Returns the size (in bytes) occupied by one parameter of type `dtype`.
+    Taken from https://github.com/huggingface/transformers/blob/74d9d0cebb0263a3f8ab9c280569170cc74651d0/src/transformers/modeling_tf_utils.py#L608.
+    NOTE: why not `tensor.numpy().nbytes`?
+    Example:
+    ```py
+    >>> _dtype_byte_size(tf.float32)
+    4
+    ```
+    """
+    import tensorflow as tf
+
+    if dtype == tf.bool:
+        return 1 / 8
+    bit_search = re.search(r"[^\d](\d+)$", dtype.name)
+    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

huggingface_hub/serialization/_torch.py

@@ -0,0 +1,199 @@
+# Copyright 2024 The HuggingFace Team. 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.
+"""Contains pytorch-specific helpers."""
+import importlib
+from functools import lru_cache
+from typing import TYPE_CHECKING, Dict, Tuple
+
+from ._base import FILENAME_PATTERN, MAX_SHARD_SIZE, StateDictSplit, split_state_dict_into_shards_factory
+
+
+if TYPE_CHECKING:
+    import torch
+
+
+def split_torch_state_dict_into_shards(
+    state_dict: Dict[str, "torch.Tensor"],
+    *,
+    filename_pattern: str = FILENAME_PATTERN,
+    max_shard_size: int = MAX_SHARD_SIZE,
+) -> StateDictSplit:
+    """
+    Split a model state dictionary in shards so that each shard is smaller than a given size.
+
+    The shards are determined by iterating through the `state_dict` in the order of its keys. There is no optimization
+    made to make each shard as close as possible to the maximum size passed. For example, if the limit is 10GB and we
+    have tensors of sizes [6GB, 6GB, 2GB, 6GB, 2GB, 2GB] they will get sharded as [6GB], [6+2GB], [6+2+2GB] and not
+    [6+2+2GB], [6+2GB], [6GB].
+
+    <Tip warning={true}>
+
+    If one of the model's tensor is bigger than `max_shard_size`, it will end up in its own shard which will have a
+    size greater than `max_shard_size`.
+
+    </Tip>
+
+    Args:
+        state_dict (`Dict[str, torch.Tensor]`):
+            The state dictionary to save.
+        filename_pattern (`str`, *optional*):
+            The pattern to generate the files names in which the model will be saved. Pattern must be a string that
+            can be formatted with `filename_pattern.format(suffix=...)` and must contain the keyword `suffix`
+            Defaults to `"model{suffix}.safetensors"`.
+        max_shard_size (`int` or `str`, *optional*):
+            The maximum size of each shard, in bytes. Defaults to 5GB.
+
+    Returns:
+        [`StateDictSplit`]: A `StateDictSplit` object containing the shards and the index to retrieve them.
+
+    Example:
+    ```py
+    >>> import json
+    >>> import os
+    >>> from safetensors.torch import save_file as safe_save_file
+    >>> from huggingface_hub import split_torch_state_dict_into_shards
+
+    >>> def save_state_dict(state_dict: Dict[str, torch.Tensor], save_directory: str):
+    ...     state_dict_split = split_torch_state_dict_into_shards(state_dict)
+    ...     for filename, tensors in state_dict_split.filename_to_tensors.values():
+    ...         shard = {tensor: state_dict[tensor] for tensor in tensors}
+    ...         safe_save_file(
+    ...             shard,
+    ...             os.path.join(save_directory, filename),
+    ...             metadata={"format": "pt"},
+    ...         )
+    ...     if state_dict_split.is_sharded:
+    ...         index = {
+    ...             "metadata": state_dict_split.metadata,
+    ...             "weight_map": state_dict_split.tensor_to_filename,
+    ...         }
+    ...         with open(os.path.join(save_directory, "model.safetensors.index.json"), "w") as f:
+    ...             f.write(json.dumps(index, indent=2))
+    ```
+    """
+    return split_state_dict_into_shards_factory(
+        state_dict,
+        max_shard_size=max_shard_size,
+        filename_pattern=filename_pattern,
+        get_tensor_size=get_tensor_size,
+        get_storage_id=get_storage_id,
+    )
+
+
+def get_storage_id(tensor: "torch.Tensor") -> Tuple["torch.device", int, int]:
+    """
+    Return unique identifier to a tensor storage.
+
+    Multiple different tensors can share the same underlying storage. For
+    example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is
+    guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with
+    non-overlapping lifetimes may have the same id.
+
+    Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/pytorch_utils.py#L278.
+    """
+    if tensor.device.type == "xla" and is_torch_tpu_available():
+        # NOTE: xla tensors dont have storage
+        # use some other unique id to distinguish.
+        # this is a XLA tensor, it must be created using torch_xla's
+        # device. So the following import is safe:
+        import torch_xla
+
+        unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor)
+    else:
+        unique_id = storage_ptr(tensor)
+
+    return tensor.device, unique_id, get_storage_size(tensor)
+
+
+def get_tensor_size(tensor: "torch.Tensor") -> int:
+    return tensor.numel() * tensor.element_size()
+
+
+@lru_cache()
+def is_torch_tpu_available(check_device=True):
+    """
+    Checks if `torch_xla` is installed and potentially if a TPU is in the environment
+
+    Taken from https://github.com/huggingface/transformers/blob/1ecf5f7c982d761b4daaa96719d162c324187c64/src/transformers/utils/import_utils.py#L463.
+    """
+    if importlib.util.find_spec("torch_xla") is not None:
+        if check_device:
+            # We need to check if `xla_device` can be found, will raise a RuntimeError if not
+            try:
+                import torch_xla.core.xla_model as xm
+
+                _ = xm.xla_device()
+                return True
+            except RuntimeError:
+                return False
+        return True
+    return False
+
+
+def storage_ptr(tensor: "torch.Tensor") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L11C1-L20C21.
+    """
+    try:
+        return tensor.untyped_storage().data_ptr()
+    except Exception:
+        # Fallback for torch==1.10
+        try:
+            return tensor.storage().data_ptr()
+        except NotImplementedError:
+            # Fallback for meta storage
+            return 0
+
+
+def get_storage_size(tensor: "torch.Tensor") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L31C1-L41C59
+    """
+    try:
+        return tensor.untyped_storage().nbytes()
+    except AttributeError:
+        # Fallback for torch==1.10
+        try:
+            return tensor.storage().size() * _get_dtype_size(tensor.dtype)
+        except NotImplementedError:
+            # Fallback for meta storage
+            # On torch >=2.0 this is the tensor size
+            return tensor.nelement() * _get_dtype_size(tensor.dtype)
+
+
+@lru_cache()
+def _get_dtype_size(dtype: "torch.dtype") -> int:
+    """
+    Taken from https://github.com/huggingface/safetensors/blob/08db34094e9e59e2f9218f2df133b7b4aaff5a99/bindings/python/py_src/safetensors/torch.py#L344
+    """
+    import torch
+
+    # torch.float8 formats require 2.1; we do not support these dtypes on earlier versions
+    _float8_e4m3fn = getattr(torch, "float8_e4m3fn", None)
+    _float8_e5m2 = getattr(torch, "float8_e5m2", None)
+    _SIZE = {
+        torch.int64: 8,
+        torch.float32: 4,
+        torch.int32: 4,
+        torch.bfloat16: 2,
+        torch.float16: 2,
+        torch.int16: 2,
+        torch.uint8: 1,
+        torch.int8: 1,
+        torch.bool: 1,
+        torch.float64: 8,
+        _float8_e4m3fn: 1,
+        _float8_e5m2: 1,
+    }
+    return _SIZE[dtype]

huggingface_hub/templates/datasetcard_template.md

@@ -140,4 +140,4 @@
 
 ## Dataset Card Contact
 
-{{ dataset_card_contact | default("[More Information Needed]", true)}}
+{{ dataset_card_contact | default("[More Information Needed]", true)}}

huggingface_hub/templates/modelcard_template.md

@@ -82,7 +82,7 @@ Use the code below to get started with the model.
 
 {{ training_data | default("[More Information Needed]", true)}}
 
-### Training Procedure 
+### Training Procedure
 
 <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
 
@@ -198,6 +198,3 @@ Carbon emissions can be estimated using the [Machine Learning Impact calculator]
 ## Model Card Contact
 
 {{ model_card_contact | default("[More Information Needed]", true)}}
-
-
-

huggingface_hub/utils/__init__.py

@@ -1,4 +1,3 @@
-# flake8: noqa
 #!/usr/bin/env python
 # coding=utf-8
 # Copyright 2021 The HuggingFace Inc. team. All rights reserved.
@@ -15,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
+# ruff: noqa: F401
+
 from . import tqdm as _tqdm  # _tqdm is the module
 from ._cache_assets import cached_assets_path
 from ._cache_manager import (
@@ -31,6 +32,7 @@ from ._chunk_utils import chunk_iterable
 from ._datetime import parse_datetime
 from ._errors import (
     BadRequestError,
+    DisabledRepoError,
     EntryNotFoundError,
     FileMetadataError,
     GatedRepoError,
@@ -40,15 +42,14 @@ from ._errors import (
     RevisionNotFoundError,
     hf_raise_for_status,
 )
-from ._token import get_token
+from ._experimental import experimental
 from ._fixes import SoftTemporaryDirectory, yaml_dump
 from ._git_credential import list_credential_helpers, set_git_credential, unset_git_credential
-from ._headers import build_hf_headers, get_token_to_send, LocalTokenNotFoundError
+from ._headers import LocalTokenNotFoundError, build_hf_headers, get_token_to_send
 from ._hf_folder import HfFolder
-from ._http import configure_http_backend, get_session, http_backoff, reset_sessions, OfflineModeIsEnabled
+from ._http import OfflineModeIsEnabled, configure_http_backend, get_session, http_backoff, reset_sessions
 from ._pagination import paginate
-from ._paths import filter_repo_objects, IGNORE_GIT_FOLDER_PATTERNS
-from ._experimental import experimental
+from ._paths import IGNORE_GIT_FOLDER_PATTERNS, filter_repo_objects
 from ._runtime import (
     dump_environment_info,
     get_aiohttp_version,
@@ -70,28 +71,32 @@ from ._runtime import (
     is_aiohttp_available,
     is_fastai_available,
     is_fastcore_available,
-    is_numpy_available,
     is_google_colab,
     is_gradio_available,
     is_graphviz_available,
     is_hf_transfer_available,
     is_jinja_available,
     is_notebook,
+    is_numpy_available,
+    is_package_available,
     is_pillow_available,
     is_pydantic_available,
     is_pydot_available,
+    is_safetensors_available,
     is_tensorboard_available,
     is_tf_available,
     is_torch_available,
 )
 from ._safetensors import (
+    NotASafetensorsRepoError,
     SafetensorsFileMetadata,
+    SafetensorsParsingError,
     SafetensorsRepoMetadata,
     TensorInfo,
-    SafetensorsParsingError,
-    NotASafetensorsRepoError,
 )
 from ._subprocess import capture_output, run_interactive_subprocess, run_subprocess
+from ._telemetry import send_telemetry
+from ._token import get_token
 from ._validators import (
     HFValidationError,
     smoothly_deprecate_use_auth_token,
@@ -105,4 +110,3 @@ from .tqdm import (
     tqdm,
     tqdm_stream_file,
 )
-from ._telemetry import send_telemetry

huggingface_hub/utils/_datetime.py

@@ -13,13 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Contains utilities to handle datetimes in Huggingface Hub."""
-from datetime import datetime, timedelta, timezone
-
-
-# Local machine offset compared to UTC.
-# Taken from https://stackoverflow.com/a/3168394.
-# `utcoffset()` returns `None` if no offset -> empty timedelta.
-UTC_OFFSET = datetime.now(timezone.utc).astimezone().utcoffset() or timedelta()
+from datetime import datetime, timezone
 
 
 def parse_datetime(date_string: str) -> datetime:
@@ -51,16 +45,15 @@ def parse_datetime(date_string: str) -> datetime:
             If `date_string` cannot be parsed.
     """
     try:
-        # Datetime ending with a Z means "UTC". Here we parse the date as local machine
-        # timezone and then move it to the appropriate UTC timezone.
+        # Datetime ending with a Z means "UTC". We parse the date and then explicitly
+        # set the timezone to UTC.
         # See https://en.wikipedia.org/wiki/ISO_8601#Coordinated_Universal_Time_(UTC)
         # Taken from https://stackoverflow.com/a/3168394.
         if len(date_string) == 30:
             # Means timezoned-timestamp with nanoseconds precision. We need to truncate the last 3 digits.
             date_string = date_string[:-4] + "Z"
         dt = datetime.strptime(date_string, "%Y-%m-%dT%H:%M:%S.%fZ")
-        dt += UTC_OFFSET  # By default, datetime is not timezoned -> move to UTC time
-        return dt.astimezone(timezone.utc)  # Set explicit timezone
+        return dt.replace(tzinfo=timezone.utc)  # Set explicit timezone
     except ValueError as e:
         raise ValueError(
             f"Cannot parse '{date_string}' as a datetime. Date string is expected to"

huggingface_hub/utils/_errors.py

@@ -154,6 +154,24 @@ class GatedRepoError(RepositoryNotFoundError):
     """
 
 
+class DisabledRepoError(HfHubHTTPError):
+    """
+    Raised when trying to access a repository that has been disabled by its author.
+
+    Example:
+
+    ```py
+    >>> from huggingface_hub import dataset_info
+    >>> dataset_info("laion/laion-art")
+    (...)
+    huggingface_hub.utils._errors.DisabledRepoError: 403 Client Error. (Request ID: Root=1-659fc3fa-3031673e0f92c71a2260dbe2;bc6f4dfb-b30a-4862-af0a-5cfe827610d8)
+
+    Cannot access repository for url https://huggingface.co/api/datasets/laion/laion-art.
+    Access to this resource is disabled.
+    ```
+    """
+
+
 class RevisionNotFoundError(HfHubHTTPError):
     """
     Raised when trying to access a hf.co URL with a valid repository but an invalid
@@ -286,6 +304,7 @@ def hf_raise_for_status(response: Response, endpoint_name: Optional[str] = None)
         response.raise_for_status()
     except HTTPError as e:
         error_code = response.headers.get("X-Error-Code")
+        error_message = response.headers.get("X-Error-Message")
 
         if error_code == "RevisionNotFound":
             message = f"{response.status_code} Client Error." + "\n\n" + f"Revision Not Found for url: {response.url}."
@@ -301,6 +320,16 @@ def hf_raise_for_status(response: Response, endpoint_name: Optional[str] = None)
             )
             raise GatedRepoError(message, response) from e
 
+        elif error_message == "Access to this resource is disabled.":
+            message = (
+                f"{response.status_code} Client Error."
+                + "\n\n"
+                + f"Cannot access repository for url {response.url}."
+                + "\n"
+                + "Access to this resource is disabled."
+            )
+            raise DisabledRepoError(message, response) from e
+
         elif error_code == "RepoNotFound" or (
             response.status_code == 401
             and response.request is not None

huggingface_hub/utils/_hf_folder.py

@@ -13,12 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Contain helper class to retrieve/store token from/to local cache."""
-import os
 import warnings
 from pathlib import Path
 from typing import Optional
 
 from .. import constants
+from ._token import get_token
 
 
 class HfFolder:
@@ -50,12 +50,8 @@ class HfFolder:
         """
         Get token or None if not existent.
 
-        Note that a token can be also provided using the `HF_TOKEN` environment variable.
-
-        Token is saved in the huggingface home folder. You can configure it by setting
-        the `HF_HOME` environment variable. Previous location was `~/.huggingface/token`.
-        If token is found in old location but not in new location, it is copied there first.
-        For more details, see https://github.com/huggingface/huggingface_hub/issues/1232.
+        This method is deprecated in favor of [`huggingface_hub.get_token`] but is kept for backward compatibility.
+        Its behavior is the same as [`huggingface_hub.get_token`].
 
         Returns:
             `str` or `None`: The token, `None` if it doesn't exist.
@@ -66,22 +62,7 @@ class HfFolder:
         except Exception:  # if not possible (e.g. PermissionError), do not raise
             pass
 
-        # 1. Is it set by environment variable ?
-        token: Optional[str] = os.environ.get("HF_TOKEN")
-        if token is None:  # Ensure backward compatibility but doesn't have priority
-            token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
-        if token is not None:
-            token = token.replace("\r", "").replace("\n", "").strip()
-        if token != "":
-            return token
-
-        # 2. Is it set in token path ?
-        try:
-            token = cls.path_token.read_text()
-            token = token.replace("\r", "").replace("\n", "").strip()
-            return token
-        except FileNotFoundError:
-            return None
+        return get_token()
 
     # TODO: deprecate when adapted in transformers/datasets/gradio
     # @_deprecate_method(version="1.0", message="Use `huggingface_hub.logout` instead.")

huggingface_hub/utils/_runtime.py

@@ -38,6 +38,7 @@ _CANDIDATES = {
     "pillow": {"Pillow"},
     "pydantic": {"pydantic"},
     "pydot": {"pydot"},
+    "safetensors": {"safetensors"},
     "tensorboard": {"tensorboardX"},
     "tensorflow": (
         "tensorflow",
@@ -69,7 +70,7 @@ def _get_version(package_name: str) -> str:
     return _package_versions.get(package_name, "N/A")
 
 
-def _is_available(package_name: str) -> bool:
+def is_package_available(package_name: str) -> bool:
     return _get_version(package_name) != "N/A"
 
 
@@ -85,7 +86,7 @@ def get_hf_hub_version() -> str:
 
 # aiohttp
 def is_aiohttp_available() -> bool:
-    return _is_available("aiohttp")
+    return is_package_available("aiohttp")
 
 
 def get_aiohttp_version() -> str:
@@ -94,7 +95,7 @@ def get_aiohttp_version() -> str:
 
 # FastAI
 def is_fastai_available() -> bool:
-    return _is_available("fastai")
+    return is_package_available("fastai")
 
 
 def get_fastai_version() -> str:
@@ -103,7 +104,7 @@ def get_fastai_version() -> str:
 
 # Fastcore
 def is_fastcore_available() -> bool:
-    return _is_available("fastcore")
+    return is_package_available("fastcore")
 
 
 def get_fastcore_version() -> str:
@@ -112,7 +113,7 @@ def get_fastcore_version() -> str:
 
 # FastAI
 def is_gradio_available() -> bool:
-    return _is_available("gradio")
+    return is_package_available("gradio")
 
 
 def get_gradio_version() -> str:
@@ -121,7 +122,7 @@ def get_gradio_version() -> str:
 
 # Graphviz
 def is_graphviz_available() -> bool:
-    return _is_available("graphviz")
+    return is_package_available("graphviz")
 
 
 def get_graphviz_version() -> str:
@@ -130,7 +131,7 @@ def get_graphviz_version() -> str:
 
 # hf_transfer
 def is_hf_transfer_available() -> bool:
-    return _is_available("hf_transfer")
+    return is_package_available("hf_transfer")
 
 
 def get_hf_transfer_version() -> str:
@@ -139,7 +140,7 @@ def get_hf_transfer_version() -> str:
 
 # Numpy
 def is_numpy_available() -> bool:
-    return _is_available("numpy")
+    return is_package_available("numpy")
 
 
 def get_numpy_version() -> str:
@@ -148,7 +149,7 @@ def get_numpy_version() -> str:
 
 # Jinja
 def is_jinja_available() -> bool:
-    return _is_available("jinja")
+    return is_package_available("jinja")
 
 
 def get_jinja_version() -> str:
@@ -157,7 +158,7 @@ def get_jinja_version() -> str:
 
 # Pillow
 def is_pillow_available() -> bool:
-    return _is_available("pillow")
+    return is_package_available("pillow")
 
 
 def get_pillow_version() -> str:
@@ -166,7 +167,7 @@ def get_pillow_version() -> str:
 
 # Pydantic
 def is_pydantic_available() -> bool:
-    if not _is_available("pydantic"):
+    if not is_package_available("pydantic"):
         return False
     # For Pydantic, we add an extra check to test whether it is correctly installed or not. If both pydantic 2.x and
     # typing_extensions<=4.5.0 are installed, then pydantic will fail at import time. This should not happen when
@@ -195,7 +196,7 @@ def get_pydantic_version() -> str:
 
 # Pydot
 def is_pydot_available() -> bool:
-    return _is_available("pydot")
+    return is_package_available("pydot")
 
 
 def get_pydot_version() -> str:
@@ -204,7 +205,7 @@ def get_pydot_version() -> str:
 
 # Tensorboard
 def is_tensorboard_available() -> bool:
-    return _is_available("tensorboard")
+    return is_package_available("tensorboard")
 
 
 def get_tensorboard_version() -> str:
@@ -213,7 +214,7 @@ def get_tensorboard_version() -> str:
 
 # Tensorflow
 def is_tf_available() -> bool:
-    return _is_available("tensorflow")
+    return is_package_available("tensorflow")
 
 
 def get_tf_version() -> str:
@@ -222,13 +223,18 @@ def get_tf_version() -> str:
 
 # Torch
 def is_torch_available() -> bool:
-    return _is_available("torch")
+    return is_package_available("torch")
 
 
 def get_torch_version() -> str:
     return _get_version("torch")
 
 
+# Safetensors
+def is_safetensors_available() -> bool:
+    return is_package_available("safetensors")
+
+
 # Shell-related helpers
 try:
     # Set to `True` if script is running in a Google Colab notebook.