transformers 5.0.0__py3-none-any.whl → 5.0.0rc0__py3-none-any.whl

transformers/masking_utils.py

@@ -1,3 +1,4 @@
+# coding=utf-8
 # Copyright 2025 HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,6 +14,7 @@
 # limitations under the License.
 import itertools
 from collections.abc import Callable
+from typing import Optional, Union
 
 import torch
 import torch.nn.functional as F
@@ -20,7 +22,7 @@ import torch.nn.functional as F
 from .cache_utils import Cache
 from .configuration_utils import PreTrainedConfig
 from .utils import is_torch_xpu_available, logging
-from .utils.generic import GeneralInterface, is_flash_attention_requested
+from .utils.generic import GeneralInterface
 from .utils.import_utils import is_torch_flex_attn_available, is_torch_greater_or_equal, is_tracing
 
 
@@ -161,7 +163,9 @@ def add_offsets_to_mask_function(mask_function: Callable, q_offset: int, kv_offs
     return inner_mask
 
 
-def prepare_padding_mask(attention_mask: torch.Tensor | None, kv_length: int, kv_offset: int) -> torch.Tensor | None:
+def prepare_padding_mask(
+    attention_mask: Optional[torch.Tensor], kv_length: int, kv_offset: int
+) -> Optional[torch.Tensor]:
     """
     From the 2D attention mask, prepare the correct padding mask to use by potentially padding it.
     """
@@ -174,10 +178,10 @@ def prepare_padding_mask(attention_mask: torch.Tensor | None, kv_length: int, kv
 
 
 def _can_skip_causal_mask_xpu(
-    padding_mask: torch.Tensor | None,
+    padding_mask: Optional[torch.Tensor],
     query_length: int,
     kv_length: int,
-    local_attention_size: int | None,
+    local_attention_size: Optional[int],
 ) -> bool:
     """
     XPU-specific logic for determining if we can skip causal mask creation.
@@ -210,11 +214,11 @@ def _can_skip_causal_mask_xpu(
 
 
 def _ignore_causal_mask_sdpa(
-    padding_mask: torch.Tensor | None,
+    padding_mask: Optional[torch.Tensor],
     query_length: int,
     kv_length: int,
     kv_offset: int,
-    local_attention_size: int | None = None,
+    local_attention_size: Optional[int] = None,
 ) -> bool:
     """
     Detects whether the causal mask can be ignored in case PyTorch's SDPA is used, rather relying on SDPA's `is_causal` argument.
@@ -253,7 +257,7 @@ def _ignore_causal_mask_sdpa(
     return False
 
 
-def _ignore_bidirectional_mask_sdpa(padding_mask: torch.Tensor | None) -> bool:
+def _ignore_bidirectional_mask_sdpa(padding_mask: Optional[torch.Tensor]) -> bool:
     """
     Detects whether the bidirectional mask can be ignored in case PyTorch's SDPA is used, i.e. when there is full
     attention with no padding.
@@ -304,14 +308,14 @@ def sdpa_mask(
     kv_length: int,
     kv_offset: int = 0,
     mask_function: Callable = causal_mask_function,
-    attention_mask: torch.Tensor | None = None,
-    local_size: int | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    local_size: Optional[int] = None,
     allow_is_causal_skip: bool = True,
     allow_is_bidirectional_skip: bool = False,
     allow_torch_fix: bool = True,
     use_vmap: bool = False,
     **kwargs,
-) -> torch.Tensor | None:
+) -> Optional[torch.Tensor]:
     """
     Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that
     the element should take part in the attention computation, and False that it should not.
@@ -471,7 +475,7 @@ def eager_mask(
     kv_length: int,
     kv_offset: int = 0,
     mask_function: Callable = causal_mask_function,
-    attention_mask: torch.Tensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
     dtype: torch.dtype = torch.float32,
     allow_is_bidirectional_skip: bool = False,
     use_vmap: bool = False,
@@ -534,7 +538,7 @@ def flash_attention_mask(
     kv_length: int,
     kv_offset: int = 0,
     mask_function: Callable = causal_mask_function,
-    attention_mask: torch.Tensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
     **kwargs,
 ):
     """
@@ -573,7 +577,7 @@ def flex_attention_mask(
     kv_length: int,
     kv_offset: int = 0,
     mask_function: Callable = causal_mask_function,
-    attention_mask: torch.Tensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
     **kwargs,
 ) -> BlockMask:
     """
@@ -640,7 +644,7 @@ class AttentionMaskInterface(GeneralInterface):
 ALL_MASK_ATTENTION_FUNCTIONS: AttentionMaskInterface = AttentionMaskInterface()
 
 
-def find_packed_sequence_indices(position_ids: torch.Tensor) -> torch.Tensor | None:
+def find_packed_sequence_indices(position_ids: torch.Tensor) -> Optional[torch.Tensor]:
     """
     Find the indices of the sequence to which each new query token in the sequence belongs when using packed
     tensor format (i.e. several sequences packed in the same batch dimension).
@@ -675,12 +679,12 @@ def find_packed_sequence_indices(position_ids: torch.Tensor) -> torch.Tensor | N
 def _preprocess_mask_arguments(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | BlockMask | None,
+    attention_mask: Optional[Union[torch.Tensor, BlockMask]],
     cache_position: torch.Tensor,
-    past_key_values: Cache | None,
-    position_ids: torch.Tensor | None,
-    layer_idx: int | None,
-) -> tuple[bool, torch.Tensor | BlockMask | None, int, int]:
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor],
+    layer_idx: Optional[int],
+) -> tuple[bool, Optional[Union[torch.Tensor, BlockMask]], int, int]:
     """
     Perform some common pre-processing of the mask arguments we get from the modeling code. Mostly determine the
     key-value length and offsets, and if we should early exit or not.
@@ -766,13 +770,13 @@ def _preprocess_mask_arguments(
 def create_causal_mask(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | None,
+    attention_mask: Optional[torch.Tensor],
     cache_position: torch.Tensor,
-    past_key_values: Cache | None,
-    position_ids: torch.Tensor | None = None,
-    or_mask_function: Callable | None = None,
-    and_mask_function: Callable | None = None,
-) -> torch.Tensor | BlockMask | None:
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor] = None,
+    or_mask_function: Optional[Callable] = None,
+    and_mask_function: Optional[Callable] = None,
+) -> Optional[Union[torch.Tensor, BlockMask]]:
     """
     Create a standard causal mask based on the attention implementation used (stored in the config). If `past_key_values`
     has an hybrid cache structure, this function will return the mask corresponding to one of the "full_attention" layers (to align
@@ -869,11 +873,11 @@ def create_causal_mask(
 def create_bidirectional_mask(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | None,
-    encoder_hidden_states: torch.Tensor | None = None,
-    or_mask_function: Callable | None = None,
-    and_mask_function: Callable | None = None,
-) -> torch.Tensor | BlockMask | None:
+    attention_mask: Optional[torch.Tensor],
+    encoder_hidden_states: Optional[torch.Tensor] = None,
+    or_mask_function: Optional[Callable] = None,
+    and_mask_function: Optional[Callable] = None,
+) -> Optional[Union[torch.Tensor, BlockMask]]:
     """
     Create a standard bidirectional mask based on the attention implementation used (stored in the config).
 
@@ -957,13 +961,13 @@ def create_bidirectional_mask(
 def create_sliding_window_causal_mask(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | None,
+    attention_mask: Optional[torch.Tensor],
     cache_position: torch.Tensor,
-    past_key_values: Cache | None,
-    position_ids: torch.Tensor | None = None,
-    or_mask_function: Callable | None = None,
-    and_mask_function: Callable | None = None,
-) -> torch.Tensor | BlockMask | None:
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor] = None,
+    or_mask_function: Optional[Callable] = None,
+    and_mask_function: Optional[Callable] = None,
+) -> Optional[Union[torch.Tensor, BlockMask]]:
     """
     Create a sliding window causal mask based on the attention implementation used (stored in the config). This type
     of attention pattern was mostly democratized by Mistral. If `past_key_values` has an hybrid cache structure, this
@@ -1061,13 +1065,13 @@ def create_sliding_window_causal_mask(
 def create_chunked_causal_mask(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | None,
+    attention_mask: Optional[torch.Tensor],
     cache_position: torch.Tensor,
-    past_key_values: Cache | None,
-    position_ids: torch.Tensor | None = None,
-    or_mask_function: Callable | None = None,
-    and_mask_function: Callable | None = None,
-) -> torch.Tensor | BlockMask | None:
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor] = None,
+    or_mask_function: Optional[Callable] = None,
+    and_mask_function: Optional[Callable] = None,
+) -> Optional[Union[torch.Tensor, BlockMask]]:
     """
     Create a chunked attention causal mask based on the attention implementation used (stored in the config). This type
     of attention pattern was mostly democratized by Llama4. If `past_key_values` has an hybrid cache structure, this
@@ -1112,10 +1116,10 @@ def create_chunked_causal_mask(
     if chunk_size is None:
         raise ValueError("Could not find an `attention_chunk_size` argument in the config, or it is not set")
 
-    # Raise if using chunked attention on context too large with FA
-    if is_flash_attention_requested(config) and kv_length + kv_offset > chunk_size:
+    # Raise if using chunked attention on context too large with FA2
+    if config._attn_implementation == "flash_attention_2" and kv_length + kv_offset > chunk_size:
         raise ValueError(
-            "Flash attention cannot handle chunked attention, and the key-value length is larger than the chunk size so the "
+            "Flash attention 2 cannot handle chunked attention, and the key-value length is larger than the chunk size so the "
             "chunked pattern cannot be respected. You should use another `attn_implementation` when instantiating the model"
         )
 
@@ -1186,12 +1190,12 @@ LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING = {
 def create_masks_for_generate(
     config: PreTrainedConfig,
     input_embeds: torch.Tensor,
-    attention_mask: torch.Tensor | None,
+    attention_mask: Optional[torch.Tensor],
     cache_position: torch.Tensor,
-    past_key_values: Cache | None,
-    position_ids: torch.Tensor | None = None,
-    or_mask_function: Callable | None = None,
-    and_mask_function: Callable | None = None,
+    past_key_values: Optional[Cache],
+    position_ids: Optional[torch.Tensor] = None,
+    or_mask_function: Optional[Callable] = None,
+    and_mask_function: Optional[Callable] = None,
     **kwargs,
 ):
     """
@@ -1371,7 +1375,7 @@ class AttentionMask(torch.Tensor):
         return self.to_string()
 
     @classmethod
-    def from_tensor(cls, tensor: torch.Tensor, style: str | None = None) -> "AttentionMask":
+    def from_tensor(cls, tensor: torch.Tensor, style: Optional[str] = None) -> "AttentionMask":
         res = cls(tensor)
         res.style = style
         return res

transformers/model_debugging_utils.py

@@ -19,6 +19,7 @@ import os
 import re
 from contextlib import contextmanager, redirect_stdout
 from io import StringIO
+from typing import Optional
 
 from .utils import logging
 from .utils.import_utils import is_torch_available, requires
@@ -67,7 +68,7 @@ def _dtensor_repr(x):
 
 
 def _serialize_tensor_like_io(
-    value, debug_path: str | None = None, use_repr: bool = True, path_to_value: str | None = None
+    value, debug_path: Optional[str] = None, use_repr: bool = True, path_to_value: Optional[str] = None
 ):
     """
     Converts Tensors and DTensors to a JSON-serializable dictionary representation.
@@ -115,7 +116,7 @@ def _serialize_tensor_like_io(
     return out
 
 
-def _serialize_io(value, debug_path: str | None = None, use_repr: bool = True, path_to_value: str | None = None):
+def _serialize_io(value, debug_path: Optional[str] = None, use_repr: bool = True, path_to_value: Optional[str] = None):
     """
     Recursively build a JSON-serializable Python structure from `value`.
     Tensors and DTensors become either sanitized repr strings, or are saved to disk as SafeTensors files and their
@@ -224,7 +225,7 @@ def prune_intermediate_layers(node):
         prune_intermediate_layers(child)
 
 
-def log_model_debug_trace(debug_path: str | None, model):
+def log_model_debug_trace(debug_path: Optional[str], model):
     if debug_path:
         try:
             os.makedirs(debug_path, exist_ok=True)
@@ -398,7 +399,7 @@ def _attach_debugger_logic(
 @contextmanager
 def model_addition_debugger_context(
     model,
-    debug_path: str | None = None,
+    debug_path: Optional[str] = None,
     do_prune_layers: bool = True,
     use_repr: bool = True,
 ):

transformers/modelcard.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 """Configuration base class and utilities."""
 
+import copy
+import json
 import os
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, Optional, Union
 
 import httpx
 import yaml
-from huggingface_hub import is_offline_mode, model_info
+from huggingface_hub import model_info
 from huggingface_hub.errors import OfflineModeIsEnabled
 from huggingface_hub.utils import HFValidationError
 
@@ -44,7 +47,10 @@ from .models.auto.modeling_auto import (
 )
 from .training_args import ParallelMode
 from .utils import (
+    MODEL_CARD_NAME,
+    cached_file,
     is_datasets_available,
+    is_offline_mode,
     is_tokenizers_available,
     is_torch_available,
     logging,
@@ -71,6 +77,179 @@ TASK_MAPPING = {
 logger = logging.get_logger(__name__)
 
 
+class ModelCard:
+    r"""
+    Structured Model Card class. Store model card as well as methods for loading/downloading/saving model cards.
+
+    Please read the following paper for details and explanation on the sections: "Model Cards for Model Reporting" by
+    Margaret Mitchell, Simone Wu, Andrew Zaldivar, Parker Barnes, Lucy Vasserman, Ben Hutchinson, Elena Spitzer,
+    Inioluwa Deborah Raji and Timnit Gebru for the proposal behind model cards. Link: https://huggingface.co/papers/1810.03993
+
+    Note: A model card can be loaded and saved to disk.
+    """
+
+    def __init__(self, **kwargs):
+        warnings.warn(
+            "The class `ModelCard` is deprecated and will be removed in version 5 of Transformers", FutureWarning
+        )
+        # Recommended attributes from https://huggingface.co/papers/1810.03993 (see papers)
+        self.model_details = kwargs.pop("model_details", {})
+        self.intended_use = kwargs.pop("intended_use", {})
+        self.factors = kwargs.pop("factors", {})
+        self.metrics = kwargs.pop("metrics", {})
+        self.evaluation_data = kwargs.pop("evaluation_data", {})
+        self.training_data = kwargs.pop("training_data", {})
+        self.quantitative_analyses = kwargs.pop("quantitative_analyses", {})
+        self.ethical_considerations = kwargs.pop("ethical_considerations", {})
+        self.caveats_and_recommendations = kwargs.pop("caveats_and_recommendations", {})
+
+        # Open additional attributes
+        for key, value in kwargs.items():
+            try:
+                setattr(self, key, value)
+            except AttributeError as err:
+                logger.error(f"Can't set {key} with value {value} for {self}")
+                raise err
+
+    def save_pretrained(self, save_directory_or_file):
+        """Save a model card object to the directory or file `save_directory_or_file`."""
+        if os.path.isdir(save_directory_or_file):
+            # If we save using the predefined names, we can load using `from_pretrained`
+            output_model_card_file = os.path.join(save_directory_or_file, MODEL_CARD_NAME)
+        else:
+            output_model_card_file = save_directory_or_file
+
+        self.to_json_file(output_model_card_file)
+        logger.info(f"Model card saved in {output_model_card_file}")
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        r"""
+        Instantiate a [`ModelCard`] from a pre-trained model model card.
+
+        Parameters:
+            pretrained_model_name_or_path: either:
+
+                - a string, the *model id* of a pretrained model card hosted inside a model repo on huggingface.co.
+                - a path to a *directory* containing a model card file saved using the [`~ModelCard.save_pretrained`]
+                  method, e.g.: `./my_model_directory/`.
+                - a path or url to a saved model card JSON *file*, e.g.: `./my_model_directory/modelcard.json`.
+
+            cache_dir: (*optional*) string:
+                Path to a directory in which a downloaded pre-trained model card should be cached if the standard cache
+                should not be used.
+
+            kwargs: (*optional*) dict: key/value pairs with which to update the ModelCard object after loading.
+
+                - The values in kwargs of any keys which are model card attributes will be used to override the loaded
+                  values.
+                - Behavior concerning key/value pairs whose keys are *not* model card attributes is controlled by the
+                  *return_unused_kwargs* keyword parameter.
+
+            proxies: (*optional*) dict, default None:
+                A dictionary of proxy servers to use by protocol or endpoint, e.g.: {'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}. The proxies are used on each request.
+
+            return_unused_kwargs: (*optional*) bool:
+
+                - If False, then this function returns just the final model card object.
+                - If True, then this functions returns a tuple *(model card, unused_kwargs)* where *unused_kwargs* is a
+                  dictionary consisting of the key/value pairs whose keys are not model card attributes: ie the part of
+                  kwargs which has not been used to update *ModelCard* and is otherwise ignored.
+
+        Examples:
+
+        ```python
+        # Download model card from huggingface.co and cache.
+        modelcard = ModelCard.from_pretrained("google-bert/bert-base-uncased")
+        # Model card was saved using *save_pretrained('./test/saved_model/')*
+        modelcard = ModelCard.from_pretrained("./test/saved_model/")
+        modelcard = ModelCard.from_pretrained("./test/saved_model/modelcard.json")
+        modelcard = ModelCard.from_pretrained("google-bert/bert-base-uncased", output_attentions=True, foo=False)
+        ```"""
+        cache_dir = kwargs.pop("cache_dir", None)
+        proxies = kwargs.pop("proxies", None)
+        return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+        from_pipeline = kwargs.pop("_from_pipeline", None)
+
+        user_agent = {"file_type": "model_card"}
+        if from_pipeline is not None:
+            user_agent["using_pipeline"] = from_pipeline
+
+        is_local = os.path.isdir(pretrained_model_name_or_path)
+        if os.path.isfile(pretrained_model_name_or_path):
+            resolved_model_card_file = pretrained_model_name_or_path
+            is_local = True
+        else:
+            try:
+                # Load from URL or cache if already cached
+                resolved_model_card_file = cached_file(
+                    pretrained_model_name_or_path,
+                    filename=MODEL_CARD_NAME,
+                    cache_dir=cache_dir,
+                    proxies=proxies,
+                    user_agent=user_agent,
+                )
+                if is_local:
+                    logger.info(f"loading model card file {resolved_model_card_file}")
+                else:
+                    logger.info(f"loading model card file {MODEL_CARD_NAME} from cache at {resolved_model_card_file}")
+                # Load model card
+                modelcard = cls.from_json_file(resolved_model_card_file)
+
+            except (OSError, json.JSONDecodeError):
+                # We fall back on creating an empty model card
+                modelcard = cls()
+
+        # Update model card with kwargs if needed
+        to_remove = []
+        for key, value in kwargs.items():
+            if hasattr(modelcard, key):
+                setattr(modelcard, key, value)
+                to_remove.append(key)
+        for key in to_remove:
+            kwargs.pop(key, None)
+
+        logger.info(f"Model card: {modelcard}")
+        if return_unused_kwargs:
+            return modelcard, kwargs
+        else:
+            return modelcard
+
+    @classmethod
+    def from_dict(cls, json_object):
+        """Constructs a `ModelCard` from a Python dictionary of parameters."""
+        return cls(**json_object)
+
+    @classmethod
+    def from_json_file(cls, json_file):
+        """Constructs a `ModelCard` from a json file of parameters."""
+        with open(json_file, encoding="utf-8") as reader:
+            text = reader.read()
+        dict_obj = json.loads(text)
+        return cls(**dict_obj)
+
+    def __eq__(self, other):
+        return self.__dict__ == other.__dict__
+
+    def __repr__(self):
+        return str(self.to_json_string())
+
+    def to_dict(self):
+        """Serializes this instance to a Python dictionary."""
+        output = copy.deepcopy(self.__dict__)
+        return output
+
+    def to_json_string(self):
+        """Serializes this instance to a JSON string."""
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+
+    def to_json_file(self, json_file_path):
+        """Save this instance to a json file."""
+        with open(json_file_path, "w", encoding="utf-8") as writer:
+            writer.write(self.to_json_string())
+
+
 AUTOGENERATED_TRAINER_COMMENT = """
 <!-- This model card has been generated automatically according to the information the Trainer had access to. You
 should probably proofread and complete it, then remove this comment. -->
@@ -174,19 +353,19 @@ def _get_mapping_values(mapping):
 @dataclass
 class TrainingSummary:
     model_name: str
-    language: str | list[str] | None = None
-    license: str | None = None
-    tags: str | list[str] | None = None
-    finetuned_from: str | None = None
-    tasks: str | list[str] | None = None
-    dataset: str | list[str] | None = None
-    dataset_tags: str | list[str] | None = None
-    dataset_args: str | list[str] | None = None
-    dataset_metadata: dict[str, Any] | None = None
-    eval_results: dict[str, float] | None = None
-    eval_lines: list[str] | None = None
-    hyperparameters: dict[str, Any] | None = None
-    source: str | None = "trainer"
+    language: Optional[Union[str, list[str]]] = None
+    license: Optional[str] = None
+    tags: Optional[Union[str, list[str]]] = None
+    finetuned_from: Optional[str] = None
+    tasks: Optional[Union[str, list[str]]] = None
+    dataset: Optional[Union[str, list[str]]] = None
+    dataset_tags: Optional[Union[str, list[str]]] = None
+    dataset_args: Optional[Union[str, list[str]]] = None
+    dataset_metadata: Optional[dict[str, Any]] = None
+    eval_results: Optional[dict[str, float]] = None
+    eval_lines: Optional[list[str]] = None
+    hyperparameters: Optional[dict[str, Any]] = None
+    source: Optional[str] = "trainer"
 
     def __post_init__(self):
         # Infer default license from the checkpoint used, if possible.

transformers/modeling_attn_mask_utils.py

@@ -18,7 +18,7 @@ and will be removed in the future.
 """
 
 from dataclasses import dataclass
-from typing import Union
+from typing import Optional, Union
 
 import torch
 
@@ -60,7 +60,7 @@ class AttentionMaskConverter:
     is_causal: bool
     sliding_window: int
 
-    def __init__(self, is_causal: bool, sliding_window: int | None = None):
+    def __init__(self, is_causal: bool, sliding_window: Optional[int] = None):
         self.is_causal = is_causal
         self.sliding_window = sliding_window
 
@@ -76,7 +76,7 @@ class AttentionMaskConverter:
         key_value_length: int,
         dtype: torch.dtype,
         device: Union[torch.device, "str"] = "cpu",
-    ) -> torch.Tensor | None:
+    ) -> Optional[torch.Tensor]:
         """
         Creates a causal 4D mask of (bsz, head_dim=1, query_length, key_value_length) shape and adds large negative
         bias to upper right hand triangular matrix (causal mask).
@@ -107,7 +107,7 @@ class AttentionMaskConverter:
         attention_mask_2d: torch.Tensor,
         query_length: int,
         dtype: torch.dtype,
-        key_value_length: int | None = None,
+        key_value_length: Optional[int] = None,
     ) -> torch.Tensor:
         """
         Converts 2D attention mask to 4D attention mask by expanding mask to (bsz, head_dim=1, query_length,
@@ -155,7 +155,7 @@ class AttentionMaskConverter:
         dtype: torch.dtype,
         device: torch.device,
         past_key_values_length: int = 0,
-        sliding_window: int | None = None,
+        sliding_window: Optional[int] = None,
     ):
         """
         Make causal mask used for bi-directional self-attention.
@@ -184,7 +184,7 @@ class AttentionMaskConverter:
         return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 
     @staticmethod
-    def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
+    def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
         """
         Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
         """
@@ -248,10 +248,10 @@ class AttentionMaskConverter:
 
     @staticmethod
     def _ignore_causal_mask_sdpa(
-        attention_mask: torch.Tensor | None,
+        attention_mask: Optional[torch.Tensor],
         inputs_embeds: torch.Tensor,
         past_key_values_length: int,
-        sliding_window: int | None = None,
+        sliding_window: Optional[int] = None,
         is_training: bool = False,
     ) -> bool:
         """
@@ -306,11 +306,11 @@ class AttentionMaskConverter:
 
 
 def _prepare_4d_causal_attention_mask(
-    attention_mask: torch.Tensor | None,
-    input_shape: torch.Size | tuple | list,
+    attention_mask: Optional[torch.Tensor],
+    input_shape: Union[torch.Size, tuple, list],
     inputs_embeds: torch.Tensor,
     past_key_values_length: int,
-    sliding_window: int | None = None,
+    sliding_window: Optional[int] = None,
 ):
     """
     Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
@@ -359,11 +359,11 @@ def _prepare_4d_causal_attention_mask(
 
 # Adapted from _prepare_4d_causal_attention_mask
 def _prepare_4d_causal_attention_mask_for_sdpa(
-    attention_mask: torch.Tensor | None,
-    input_shape: torch.Size | tuple | list,
+    attention_mask: Optional[torch.Tensor],
+    input_shape: Union[torch.Size, tuple, list],
     inputs_embeds: torch.Tensor,
     past_key_values_length: int,
-    sliding_window: int | None = None,
+    sliding_window: Optional[int] = None,
 ):
     """
     Prepares the correct `attn_mask` argument to be used by `torch.nn.functional.scaled_dot_product_attention`.
@@ -416,7 +416,7 @@ def _prepare_4d_causal_attention_mask_for_sdpa(
     return expanded_4d_mask
 
 
-def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
+def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
     """
     Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
     `(batch_size, key_value_length)`
@@ -432,7 +432,7 @@ def _prepare_4d_attention_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len:
     return AttentionMaskConverter._expand_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
 
 
-def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None):
+def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
     """
     Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
     `(batch_size, key_value_length)`
@@ -456,12 +456,12 @@ def _prepare_4d_attention_mask_for_sdpa(mask: torch.Tensor, dtype: torch.dtype,
 
 
 def _create_4d_causal_attention_mask(
-    input_shape: torch.Size | tuple | list,
+    input_shape: Union[torch.Size, tuple, list],
     dtype: torch.dtype,
     device: torch.device,
     past_key_values_length: int = 0,
-    sliding_window: int | None = None,
-) -> torch.Tensor | None:
+    sliding_window: Optional[int] = None,
+) -> Optional[torch.Tensor]:
     """
     Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)`