Voice et bot modif

venv/lib/python3.12/site-packages/huggingface_hub/serialization/_base.py

@@ -13,15 +13,16 @@
 # limitations under the License.
 """Contains helpers to split tensors into shards."""
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import Any, Callable, Optional, TypeVar, Union
+from typing import Any, TypeVar
 
 from .. import logging
 
 
 TensorT = TypeVar("TensorT")
 TensorSizeFn_T = Callable[[TensorT], int]
-StorageIDFn_T = Callable[[TensorT], Optional[Any]]
+StorageIDFn_T = Callable[[TensorT], Any | None]
 
 MAX_SHARD_SIZE = "5GB"
 SIZE_UNITS = {
@@ -52,7 +53,7 @@ def split_state_dict_into_shards_factory(
     get_storage_size: TensorSizeFn_T,
     filename_pattern: str,
     get_storage_id: StorageIDFn_T = lambda tensor: None,
-    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
+    max_shard_size: int | str = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.

venv/lib/python3.12/site-packages/huggingface_hub/serialization/_dduf.py

@@ -4,10 +4,11 @@ import mmap
 import os
 import shutil
 import zipfile
+from collections.abc import Generator, Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Generator, Iterable, Union
+from typing import Any
 
 from ..errors import DDUFCorruptedFileError, DDUFExportError, DDUFInvalidEntryNameError
 
@@ -87,7 +88,7 @@ class DDUFEntry:
             return f.read(self.length).decode(encoding=encoding)
 
 
-def read_dduf_file(dduf_path: Union[os.PathLike, str]) -> dict[str, DDUFEntry]:
+def read_dduf_file(dduf_path: os.PathLike | str) -> dict[str, DDUFEntry]:
     """
     Read a DDUF file and return a dictionary of entries.
 
@@ -156,9 +157,7 @@ def read_dduf_file(dduf_path: Union[os.PathLike, str]) -> dict[str, DDUFEntry]:
     return entries
 
 
-def export_entries_as_dduf(
-    dduf_path: Union[str, os.PathLike], entries: Iterable[tuple[str, Union[str, Path, bytes]]]
-) -> None:
+def export_entries_as_dduf(dduf_path: str | os.PathLike, entries: Iterable[tuple[str, str | Path | bytes]]) -> None:
     """Write a DDUF file from an iterable of entries.
 
     This is a lower-level helper than [`export_folder_as_dduf`] that allows more flexibility when serializing data.
@@ -247,7 +246,7 @@ def export_entries_as_dduf(
     logger.info(f"Done writing DDUF file {dduf_path}")
 
 
-def export_folder_as_dduf(dduf_path: Union[str, os.PathLike], folder_path: Union[str, os.PathLike]) -> None:
+def export_folder_as_dduf(dduf_path: str | os.PathLike, folder_path: str | os.PathLike) -> None:
     """
     Export a folder as a DDUF file.
 
@@ -283,7 +282,7 @@ def export_folder_as_dduf(dduf_path: Union[str, os.PathLike], folder_path: Union
     export_entries_as_dduf(dduf_path, _iterate_over_folder())
 
 
-def _dump_content_in_archive(archive: zipfile.ZipFile, filename: str, content: Union[str, os.PathLike, bytes]) -> None:
+def _dump_content_in_archive(archive: zipfile.ZipFile, filename: str, content: str | os.PathLike | bytes) -> None:
     with archive.open(filename, "w", force_zip64=True) as archive_fh:
         if isinstance(content, (str, Path)):
             content_path = Path(content)
@@ -295,7 +294,7 @@ def _dump_content_in_archive(archive: zipfile.ZipFile, filename: str, content: U
             raise DDUFExportError(f"Invalid content type for {filename}. Must be str, Path or bytes.")
 
 
-def _load_content(content: Union[str, Path, bytes]) -> bytes:
+def _load_content(content: str | Path | bytes) -> bytes:
     """Load the content of an entry as bytes.
 
     Used only for small checks (not to dump content into archive).

venv/lib/python3.12/site-packages/huggingface_hub/serialization/_torch.py

@@ -14,13 +14,15 @@
 """Contains pytorch-specific helpers."""
 
 import importlib
+import importlib.util
 import json
 import os
 import re
 from collections import defaultdict, namedtuple
+from collections.abc import Iterable
 from functools import lru_cache
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Optional, Union
+from pathlib import Path, PureWindowsPath
+from typing import TYPE_CHECKING, Any, NamedTuple, Union
 
 from packaging import version
 
@@ -38,15 +40,15 @@ if TYPE_CHECKING:
 
 def save_torch_model(
     model: "torch.nn.Module",
-    save_directory: Union[str, Path],
+    save_directory: str | Path,
     *,
-    filename_pattern: Optional[str] = None,
+    filename_pattern: str | None = None,
     force_contiguous: bool = True,
-    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
-    metadata: Optional[dict[str, str]] = None,
+    max_shard_size: int | str = MAX_SHARD_SIZE,
+    metadata: dict[str, str] | None = None,
     safe_serialization: bool = True,
     is_main_process: bool = True,
-    shared_tensors_to_discard: Optional[list[str]] = None,
+    shared_tensors_to_discard: list[str] | None = None,
 ):
     """
     Saves a given torch model to disk, handling sharding and shared tensors issues.
@@ -132,15 +134,15 @@ def save_torch_model(
 
 def save_torch_state_dict(
     state_dict: dict[str, "torch.Tensor"],
-    save_directory: Union[str, Path],
+    save_directory: str | Path,
     *,
-    filename_pattern: Optional[str] = None,
+    filename_pattern: str | None = None,
     force_contiguous: bool = True,
-    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
-    metadata: Optional[dict[str, str]] = None,
+    max_shard_size: int | str = MAX_SHARD_SIZE,
+    metadata: dict[str, str] | None = None,
     safe_serialization: bool = True,
     is_main_process: bool = True,
-    shared_tensors_to_discard: Optional[list[str]] = None,
+    shared_tensors_to_discard: list[str] | None = None,
 ) -> None:
     """
     Save a model state dictionary to the disk, handling sharding and shared tensors issues.
@@ -291,7 +293,7 @@ def split_torch_state_dict_into_shards(
     state_dict: dict[str, "torch.Tensor"],
     *,
     filename_pattern: str = constants.SAFETENSORS_WEIGHTS_FILE_PATTERN,
-    max_shard_size: Union[int, str] = MAX_SHARD_SIZE,
+    max_shard_size: int | str = MAX_SHARD_SIZE,
 ) -> StateDictSplit:
     """
     Split a model state dictionary in shards so that each shard is smaller than a given size.
@@ -362,14 +364,14 @@ def split_torch_state_dict_into_shards(
 
 def load_torch_model(
     model: "torch.nn.Module",
-    checkpoint_path: Union[str, os.PathLike],
+    checkpoint_path: str | os.PathLike,
     *,
     strict: bool = False,
     safe: bool = True,
     weights_only: bool = False,
-    map_location: Optional[Union[str, "torch.device"]] = None,
+    map_location: Union[str, "torch.device"] | None = None,
     mmap: bool = False,
-    filename_pattern: Optional[str] = None,
+    filename_pattern: str | None = None,
 ) -> NamedTuple:
     """
     Load a checkpoint into a model, handling both sharded and non-sharded checkpoints.
@@ -505,17 +507,46 @@ def _load_sharded_checkpoint(
     # The index file contains mapping of parameter names to shard files
     index_path = filename_pattern.format(suffix="") + ".index.json"
     index_file = os.path.join(save_directory, index_path)
-    with open(index_file, "r", encoding="utf-8") as f:
+    with open(index_file, encoding="utf-8") as f:
         index = json.load(f)
 
-    # 2. Validate keys if in strict mode
+    # 2. Validate shard filenames from the index
+    # This prevents path traversal attacks and extension confusion attacks
+    # (e.g. a safetensors index referencing .bin pickle files)
+    expected_extension = Path(filename_pattern.format(suffix="")).suffix  # e.g. ".safetensors"
+    shard_files = list(set(index["weight_map"].values()))
+    for shard_file in shard_files:
+        # Reject anything that could escape `save_directory` on any host OS:
+        # POSIX absolute ("/tmp/x"), Windows drive ("C:x", "C:\\x"), UNC
+        # ("\\\\server\\share\\x"), rooted-without-drive ("\\x", "/x"), or
+        # ".." traversal — including "..\\x" which `os.path.isabs` never caught on POSIX.
+        #
+        # We parse with `PureWindowsPath` *regardless of host OS*: it treats both "/" and
+        # "\\" as separators and exposes `drive` / `root`, so a single check rejects a
+        # malicious index file on Linux too (e.g. if it's later opened on Windows). The
+        # only over-strict case is a POSIX filename like "a:foo" which would be parsed as
+        # drive "a:" — such names are never produced for safetensors shards and would
+        # break on Windows anyway, so rejecting them is fine.
+        win_path = PureWindowsPath(shard_file)
+        if win_path.drive or win_path.root or ".." in win_path.parts:
+            raise ValueError(
+                f"Invalid shard filename '{shard_file}' in index file '{index_file}'. "
+                "Shard filenames must be relative paths without '..' components."
+            )
+        # Reject extension mismatch (e.g. .bin shard in a .safetensors index)
+        if not shard_file.endswith(expected_extension):
+            raise ValueError(
+                f"Invalid shard filename '{shard_file}' in index file '{index_file}'. "
+                f"Expected '{expected_extension}' extension to match the index format."
+            )
+
+    # 3. Validate keys if in strict mode
     # This is done before loading any shards to fail fast
     if strict:
         _validate_keys_for_strict_loading(model, index["weight_map"].keys())
 
-    # 3. Load each shard using `load_state_dict`
+    # 4. Load each shard using `load_state_dict`
     # Get unique shard files (multiple parameters can be in same shard)
-    shard_files = list(set(index["weight_map"].values()))
     for shard_file in shard_files:
         # Load shard into memory
         shard_path = os.path.join(save_directory, shard_file)
@@ -529,7 +560,7 @@ def _load_sharded_checkpoint(
         # Explicitly remove the state dict from memory
         del state_dict
 
-    # 4. Return compatibility info
+    # 5. Return compatibility info
     loaded_keys = set(index["weight_map"].keys())
     model_keys = set(model.state_dict().keys())
     return _IncompatibleKeys(
@@ -538,11 +569,11 @@ def _load_sharded_checkpoint(
 
 
 def load_state_dict_from_file(
-    checkpoint_file: Union[str, os.PathLike],
-    map_location: Optional[Union[str, "torch.device"]] = None,
+    checkpoint_file: str | os.PathLike,
+    map_location: Union[str, "torch.device"] | None = None,
     weights_only: bool = False,
     mmap: bool = False,
-) -> Union[dict[str, "torch.Tensor"], Any]:
+) -> dict[str, "torch.Tensor"] | Any:
     """
     Loads a checkpoint file, handling both safetensors and pickle checkpoint formats.
 
@@ -682,7 +713,7 @@ def _validate_keys_for_strict_loading(
         raise RuntimeError(error_message)
 
 
-def _get_unique_id(tensor: "torch.Tensor") -> Union[int, tuple[Any, ...]]:
+def _get_unique_id(tensor: "torch.Tensor") -> int | tuple[Any, ...]:
     """Returns a unique id for plain tensor
     or a (potentially nested) Tuple of unique id for the flattened Tensor
     if the input is a wrapper tensor subclass Tensor
@@ -723,7 +754,7 @@ def _get_unique_id(tensor: "torch.Tensor") -> Union[int, tuple[Any, ...]]:
     return unique_id
 
 
-def get_torch_storage_id(tensor: "torch.Tensor") -> Optional[tuple["torch.device", Union[int, tuple[Any, ...]], int]]:
+def get_torch_storage_id(tensor: "torch.Tensor") -> tuple["torch.device", int | tuple[Any, ...], int] | None:
     """
     Return unique identifier to a tensor storage.
 
@@ -776,7 +807,7 @@ def get_torch_storage_size(tensor: "torch.Tensor") -> int:
             return tensor.nelement() * _get_dtype_size(tensor.dtype)
 
 
-@lru_cache()
+@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
@@ -797,7 +828,7 @@ def is_torch_tpu_available(check_device=True):
     return False
 
 
-def storage_ptr(tensor: "torch.Tensor") -> Union[int, tuple[Any, ...]]:
+def storage_ptr(tensor: "torch.Tensor") -> int | tuple[Any, ...]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L11.
     """
@@ -826,7 +857,7 @@ def _clean_state_dict_for_safetensors(
     state_dict: dict[str, "torch.Tensor"],
     metadata: dict[str, str],
     force_contiguous: bool = True,
-    shared_tensors_to_discard: Optional[list[str]] = None,
+    shared_tensors_to_discard: list[str] | None = None,
 ):
     """Remove shared tensors from state_dict and update metadata accordingly (for reloading).
 
@@ -927,8 +958,8 @@ def _is_complete(tensor: "torch.Tensor") -> bool:
 def _remove_duplicate_names(
     state_dict: dict[str, "torch.Tensor"],
     *,
-    preferred_names: Optional[list[str]] = None,
-    discard_names: Optional[list[str]] = None,
+    preferred_names: list[str] | None = None,
+    discard_names: list[str] | None = None,
 ) -> dict[str, list[str]]:
     """
     Taken from https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/py_src/safetensors/torch.py#L80
@@ -943,7 +974,7 @@ def _remove_duplicate_names(
     shareds = _find_shared_tensors(state_dict)
     to_remove = defaultdict(list)
     for shared in shareds:
-        complete_names = set([name for name in shared if _is_complete(state_dict[name])])
+        complete_names = {name for name in shared if _is_complete(state_dict[name])}
         if not complete_names:
             raise RuntimeError(
                 "Error while trying to find names to remove to save state dict, but found no suitable name to keep"
@@ -973,7 +1004,7 @@ def _remove_duplicate_names(
     return to_remove
 
 
-@lru_cache()
+@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