Replace Optional and Union typing with | in some source files (#42372)

Signed-off-by: Yuanyuan Chen <[email protected]>
Co-authored-by: Matt <[email protected]>

Author: cyyever

src/transformers/integrations/accelerate.py

@@ -22,7 +22,7 @@
 import re
 from collections import OrderedDict, defaultdict
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from safetensors import safe_open
 from safetensors.torch import save_file
@@ -550,14 +550,14 @@ def offload_weight(weight: torch.Tensor, weight_name: str, offload_folder: str |
 
 def _init_infer_auto_device_map(
     model: nn.Module,
-    max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None,
-    no_split_module_classes: Optional[list[str]] = None,
-    tied_parameters: Optional[list[list[str]]] = None,
+    max_memory: dict[int | str, int | str] | None = None,
+    no_split_module_classes: list[str] | None = None,
+    tied_parameters: list[list[str]] | None = None,
     hf_quantizer: "HfQuantizer | None" = None,
 ) -> tuple[
-    list[Union[int, str]],
-    dict[Union[int, str], Union[int, str]],
-    list[Union[int, str]],
+    list[int | str],
+    dict[int | str, int | str],
+    list[int | str],
     list[int],
     dict[str, int],
     list[list[str]],
@@ -620,12 +620,12 @@ def _init_infer_auto_device_map(
 
 def infer_auto_device_map(
     model: nn.Module,
-    max_memory: Optional[dict[Union[int, str], Union[int, str]]] = None,
-    no_split_module_classes: Optional[list[str]] = None,
+    max_memory: dict[int | str, int | str] | None = None,
+    no_split_module_classes: list[str] | None = None,
     verbose: bool = False,
     clean_result: bool = True,
     offload_buffers: bool = False,
-    tied_parameters: Optional[list[list[str]]] = None,
+    tied_parameters: list[list[str]] | None = None,
     hf_quantizer: "HfQuantizer | None" = None,
 ):
     """

src/transformers/integrations/bitsandbytes.py

@@ -1,7 +1,6 @@
 import inspect
 from collections import defaultdict
 from inspect import signature
-from typing import Optional
 
 from ..core_model_loading import ConversionOps
 from ..quantizers.quantizers_utils import get_module_from_name
@@ -38,7 +37,7 @@ def __init__(self, hf_quantizer):
     def convert(
         self,
         input_dict: dict[str, list[torch.Tensor]],
-        model: Optional[torch.nn.Module] = None,
+        model: torch.nn.Module | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:
@@ -95,7 +94,7 @@ def __init__(self, hf_quantizer):
     def convert(
         self,
         input_dict: dict[str, list[torch.Tensor]],
-        model: Optional[torch.nn.Module] = None,
+        model: torch.nn.Module | None = None,
         missing_keys=None,
         **kwargs,
     ) -> dict[str, torch.Tensor]:

src/transformers/integrations/finegrained_fp8.py

@@ -15,7 +15,7 @@
 
 import re
 from collections.abc import Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 from ..core_model_loading import ConversionOps
 from ..utils import is_accelerate_available, is_torch_accelerator_available, is_torch_available, logging
@@ -655,13 +655,13 @@ def convert(self, input_dict: torch.Tensor, **kwargs) -> dict[str, torch.Tensor]
 class Fp8Dequantize(ConversionOps):
     """Inverse operation of :class:`Fp8Quantize`. Takes a pair (weight, scale) and reconstructs the fp32 tensor."""
 
-    def __init__(self, block_size: Optional[tuple[int, int]] = None):
+    def __init__(self, block_size: tuple[int, int] | None = None):
         self.block_size = block_size
         self.reverse_op = Fp8Quantize
 
     def convert(
         self,
-        value: Union[Sequence[torch.Tensor], dict[str, torch.Tensor]],
+        value: Sequence[torch.Tensor] | dict[str, torch.Tensor],
         *,
         context: dict[str, Any],
     ) -> torch.Tensor:

src/transformers/integrations/tensor_parallel.py

@@ -18,7 +18,6 @@
 import os
 import re
 from functools import partial, reduce
-from typing import Optional
 
 from ..utils.import_utils import is_torch_available
 
@@ -322,7 +321,7 @@ def repack_weights(
     return final_ordered_tensor
 
 
-def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: Optional[int] = None):
+def get_tensor_shard(param, empty_param, device_mesh, rank, dim, tensor_idx: int | None = None):
     """
     Generalized tensor sharding across a multi-dimensional device mesh.
     Extract only the fraction of the parameter owned by the given `rank` when the parameter would have gone sharding at provided `dim`.

src/transformers/optimization.py

@@ -16,7 +16,6 @@
 import math
 import warnings
 from functools import partial
-from typing import Optional, Union
 
 import torch
 from torch.optim import Optimizer
@@ -283,7 +282,7 @@ def get_polynomial_decay_schedule_with_warmup(
     return LambdaLR(optimizer, lr_lambda, last_epoch)
 
 
-def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: Optional[int] = None):
+def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps: int, timescale: int | None = None):
     if current_step < num_warmup_steps:
         return float(current_step) / float(max(1, num_warmup_steps))
     shift = timescale - num_warmup_steps
@@ -292,7 +291,7 @@ def _get_inverse_sqrt_schedule_lr_lambda(current_step: int, *, num_warmup_steps:
 
 
 def get_inverse_sqrt_schedule(
-    optimizer: Optimizer, num_warmup_steps: int, timescale: Optional[int] = None, last_epoch: int = -1
+    optimizer: Optimizer, num_warmup_steps: int, timescale: int | None = None, last_epoch: int = -1
 ):
     """
     Create a schedule with an inverse square-root learning rate, from the initial lr set in the optimizer, after a
@@ -338,8 +337,8 @@ def get_cosine_with_min_lr_schedule_with_warmup(
     num_training_steps: int,
     num_cycles: float = 0.5,
     last_epoch: int = -1,
-    min_lr: Optional[float] = None,
-    min_lr_rate: Optional[float] = None,
+    min_lr: float | None = None,
+    min_lr_rate: float | None = None,
 ):
     """
     Create a schedule with a learning rate that decreases following the values of the cosine function between the
@@ -391,7 +390,7 @@ def _get_cosine_with_min_lr_schedule_with_warmup_lr_rate_lambda(
     num_training_steps: int,
     num_cycles: float,
     min_lr_rate: float = 0.0,
-    warmup_lr_rate: Optional[float] = None,
+    warmup_lr_rate: float | None = None,
 ):
     current_step = float(current_step)
     num_warmup_steps = float(num_warmup_steps)
@@ -415,9 +414,9 @@ def get_cosine_with_min_lr_schedule_with_warmup_lr_rate(
     num_training_steps: int,
     num_cycles: float = 0.5,
     last_epoch: int = -1,
-    min_lr: Optional[float] = None,
-    min_lr_rate: Optional[float] = None,
-    warmup_lr_rate: Optional[float] = None,
+    min_lr: float | None = None,
+    min_lr_rate: float | None = None,
+    warmup_lr_rate: float | None = None,
 ):
     """
     Create a schedule with a learning rate that decreases following the values of the cosine function between the
@@ -507,8 +506,8 @@ def get_wsd_schedule(
     optimizer: Optimizer,
     num_warmup_steps: int,
     num_decay_steps: int,
-    num_training_steps: Optional[int] = None,
-    num_stable_steps: Optional[int] = None,
+    num_training_steps: int | None = None,
+    num_stable_steps: int | None = None,
     warmup_type: str = "linear",
     decay_type: str = "cosine",
     min_lr_ratio: float = 0,
@@ -592,11 +591,11 @@ def get_wsd_schedule(
 
 
 def get_scheduler(
-    name: Union[str, SchedulerType],
+    name: str | SchedulerType,
     optimizer: Optimizer,
-    num_warmup_steps: Optional[int] = None,
-    num_training_steps: Optional[int] = None,
-    scheduler_specific_kwargs: Optional[dict] = None,
+    num_warmup_steps: int | None = None,
+    num_training_steps: int | None = None,
+    scheduler_specific_kwargs: dict | None = None,
 ):
     """
     Unified API to get any scheduler from its name.

src/transformers/pipelines/base.py

@@ -177,8 +177,8 @@ def inner(items):
 def load_model(
     model,
     config: AutoConfig,
-    model_classes: Optional[tuple[type, ...]] = None,
-    task: Optional[str] = None,
+    model_classes: tuple[type, ...] | None = None,
+    task: str | None = None,
     **model_kwargs,
 ):
     """
@@ -270,7 +270,7 @@ def load_model(
     return model
 
 
-def get_default_model_and_revision(targeted_task: dict, task_options: Optional[Any]) -> tuple[str, str]:
+def get_default_model_and_revision(targeted_task: dict, task_options: Any | None) -> tuple[str, str]:
     """
     Select a default model to use for a given task.
 
@@ -305,9 +305,9 @@ def get_default_model_and_revision(targeted_task: dict, task_options: Optional[A
 
 def load_assistant_model(
     model: "PreTrainedModel",
-    assistant_model: Optional[Union[str, "PreTrainedModel"]],
-    assistant_tokenizer: Optional[PreTrainedTokenizer],
-) -> tuple[Optional["PreTrainedModel"], Optional[PreTrainedTokenizer]]:
+    assistant_model: Union[str, "PreTrainedModel"] | None,
+    assistant_tokenizer: PreTrainedTokenizer | None,
+) -> tuple[Optional["PreTrainedModel"], PreTrainedTokenizer | None]:
     """
     Prepares the assistant model and the assistant tokenizer for a pipeline whose model that can call `generate`.
 
@@ -404,9 +404,9 @@ class PipelineDataFormat:
 
     def __init__(
         self,
-        output_path: Optional[str],
-        input_path: Optional[str],
-        column: Optional[str],
+        output_path: str | None,
+        input_path: str | None,
+        column: str | None,
         overwrite: bool = False,
     ):
         self.output_path = output_path
@@ -430,7 +430,7 @@ def __iter__(self):
         raise NotImplementedError()
 
     @abstractmethod
-    def save(self, data: Union[dict, list[dict]]):
+    def save(self, data: dict | list[dict]):
         """
         Save the provided data object with the representation for the current [`~pipelines.PipelineDataFormat`].
 
@@ -439,7 +439,7 @@ def save(self, data: Union[dict, list[dict]]):
         """
         raise NotImplementedError()
 
-    def save_binary(self, data: Union[dict, list[dict]]) -> str:
+    def save_binary(self, data: dict | list[dict]) -> str:
         """
         Save the provided data object as a pickle-formatted binary data on the disk.
 
@@ -460,9 +460,9 @@ def save_binary(self, data: Union[dict, list[dict]]) -> str:
     @staticmethod
     def from_str(
         format: str,
-        output_path: Optional[str],
-        input_path: Optional[str],
-        column: Optional[str],
+        output_path: str | None,
+        input_path: str | None,
+        column: str | None,
         overwrite=False,
     ) -> "PipelineDataFormat":
         """
@@ -507,9 +507,9 @@ class CsvPipelineDataFormat(PipelineDataFormat):
 
     def __init__(
         self,
-        output_path: Optional[str],
-        input_path: Optional[str],
-        column: Optional[str],
+        output_path: str | None,
+        input_path: str | None,
+        column: str | None,
         overwrite=False,
     ):
         super().__init__(output_path, input_path, column, overwrite=overwrite)
@@ -551,9 +551,9 @@ class JsonPipelineDataFormat(PipelineDataFormat):
 
     def __init__(
         self,
-        output_path: Optional[str],
-        input_path: Optional[str],
-        column: Optional[str],
+        output_path: str | None,
+        input_path: str | None,
+        column: str | None,
         overwrite=False,
     ):
         super().__init__(output_path, input_path, column, overwrite=overwrite)
@@ -617,7 +617,7 @@ def save(self, data: dict):
         """
         print(data)
 
-    def save_binary(self, data: Union[dict, list[dict]]) -> str:
+    def save_binary(self, data: dict | list[dict]) -> str:
         if self.output_path is None:
             raise KeyError(
                 "When using piped input on pipeline outputting large object requires an output file path. "
@@ -776,13 +776,13 @@ class Pipeline(_ScikitCompat, PushToHubMixin):
     def __init__(
         self,
         model: "PreTrainedModel",
-        tokenizer: Optional[PreTrainedTokenizer] = None,
+        tokenizer: PreTrainedTokenizer | None = None,
         feature_extractor: Optional[PreTrainedFeatureExtractor] = None,
-        image_processor: Optional[BaseImageProcessor] = None,
-        processor: Optional[ProcessorMixin] = None,
-        modelcard: Optional[ModelCard] = None,
+        image_processor: BaseImageProcessor | None = None,
+        processor: ProcessorMixin | None = None,
+        modelcard: ModelCard | None = None,
         task: str = "",
-        device: Optional[Union[int, "torch.device"]] = None,
+        device: Union[int, "torch.device"] | None = None,
         binary_output: bool = False,
         **kwargs,
     ):
@@ -939,7 +939,7 @@ def __init__(
 
     def save_pretrained(
         self,
-        save_directory: Union[str, os.PathLike],
+        save_directory: str | os.PathLike,
         safe_serialization: bool = True,
         **kwargs: Any,
     ):
@@ -1085,7 +1085,7 @@ def _ensure_tensor_on_device(self, inputs, device):
         else:
             return inputs
 
-    def check_model_type(self, supported_models: Union[list[str], dict]):
+    def check_model_type(self, supported_models: list[str] | dict):
         """
         Check if the model class is in supported by the pipeline.
 
@@ -1348,9 +1348,9 @@ def register_pipeline(
         self,
         task: str,
         pipeline_class: type,
-        pt_model: Optional[Union[type, tuple[type]]] = None,
-        default: Optional[dict] = None,
-        type: Optional[str] = None,
+        pt_model: type | tuple[type] | None = None,
+        default: dict | None = None,
+        type: str | None = None,
     ) -> None:
         if task in self.supported_tasks:
             logger.warning(f"{task} is already registered. Overwriting pipeline for task {task}...")